hoveyc commited on
Commit
e30bd77
·
verified ·
1 Parent(s): 618244f

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +3 -0
  2. lipsync-models/GPT_SoVITS/AR/__init__.py +0 -0
  3. lipsync-models/GPT_SoVITS/AR/data/__init__.py +0 -0
  4. lipsync-models/GPT_SoVITS/AR/data/bucket_sampler.py +162 -0
  5. lipsync-models/GPT_SoVITS/AR/data/data_module.py +75 -0
  6. lipsync-models/GPT_SoVITS/AR/data/dataset.py +320 -0
  7. lipsync-models/GPT_SoVITS/AR/models/__init__.py +0 -0
  8. lipsync-models/GPT_SoVITS/AR/models/t2s_lightning_module.py +140 -0
  9. lipsync-models/GPT_SoVITS/AR/models/t2s_model.py +361 -0
  10. lipsync-models/GPT_SoVITS/AR/models/utils.py +243 -0
  11. lipsync-models/GPT_SoVITS/AR/modules/__init__.py +0 -0
  12. lipsync-models/GPT_SoVITS/AR/modules/activation.py +426 -0
  13. lipsync-models/GPT_SoVITS/AR/modules/embedding.py +76 -0
  14. lipsync-models/GPT_SoVITS/AR/modules/lr_schedulers.py +82 -0
  15. lipsync-models/GPT_SoVITS/AR/modules/optim.py +622 -0
  16. lipsync-models/GPT_SoVITS/AR/modules/patched_mha_with_cache.py +472 -0
  17. lipsync-models/GPT_SoVITS/AR/modules/scaling.py +335 -0
  18. lipsync-models/GPT_SoVITS/AR/modules/transformer.py +374 -0
  19. lipsync-models/GPT_SoVITS/AR/text_processing/__init__.py +0 -0
  20. lipsync-models/GPT_SoVITS/AR/text_processing/phonemizer.py +78 -0
  21. lipsync-models/GPT_SoVITS/AR/text_processing/symbols.py +9 -0
  22. lipsync-models/GPT_SoVITS/AR/utils/__init__.py +37 -0
  23. lipsync-models/GPT_SoVITS/AR/utils/initialize.py +38 -0
  24. lipsync-models/GPT_SoVITS/AR/utils/io.py +34 -0
  25. lipsync-models/GPT_SoVITS/configs/s1.yaml +31 -0
  26. lipsync-models/GPT_SoVITS/configs/s1big.yaml +31 -0
  27. lipsync-models/GPT_SoVITS/configs/s1big2.yaml +31 -0
  28. lipsync-models/GPT_SoVITS/configs/s1longer.yaml +31 -0
  29. lipsync-models/GPT_SoVITS/configs/s1mq.yaml +77 -0
  30. lipsync-models/GPT_SoVITS/configs/s2.json +90 -0
  31. lipsync-models/GPT_SoVITS/configs/train.yaml +32 -0
  32. lipsync-models/GPT_SoVITS/feature_extractor/__init__.py +6 -0
  33. lipsync-models/GPT_SoVITS/feature_extractor/cnhubert.py +104 -0
  34. lipsync-models/GPT_SoVITS/feature_extractor/whisper_enc.py +25 -0
  35. lipsync-models/GPT_SoVITS/gpt_sovits.py +445 -0
  36. lipsync-models/GPT_SoVITS/inference_gui.py +340 -0
  37. lipsync-models/GPT_SoVITS/inference_webui.py +679 -0
  38. lipsync-models/GPT_SoVITS/module/__init__.py +0 -0
  39. lipsync-models/GPT_SoVITS/module/attentions.py +709 -0
  40. lipsync-models/GPT_SoVITS/module/commons.py +189 -0
  41. lipsync-models/GPT_SoVITS/module/core_vq.py +383 -0
  42. lipsync-models/GPT_SoVITS/module/data_utils.py +332 -0
  43. lipsync-models/GPT_SoVITS/module/losses.py +73 -0
  44. lipsync-models/GPT_SoVITS/module/mel_processing.py +153 -0
  45. lipsync-models/GPT_SoVITS/module/models.py +993 -0
  46. lipsync-models/GPT_SoVITS/module/modules.py +923 -0
  47. lipsync-models/GPT_SoVITS/module/mrte_model.py +192 -0
  48. lipsync-models/GPT_SoVITS/module/quantize.py +119 -0
  49. lipsync-models/GPT_SoVITS/module/transforms.py +209 -0
  50. lipsync-models/GPT_SoVITS/prepare_datasets/1-get-text.py +131 -0
.gitattributes CHANGED
@@ -38,3 +38,6 @@ unet/flux1-dev.sft filter=lfs diff=lfs merge=lfs -text
38
  vae/ae.sft filter=lfs diff=lfs merge=lfs -text
39
  cache-huggingface/hub/models--QuanSun--EVA-CLIP/blobs/84c3a17a228c567a155259b2245b0b59072bf7da510260a0a02ec54de6d50b05 filter=lfs diff=lfs merge=lfs -text
40
  cache-huggingface/hub/models--bert-base-uncased/blobs/68d45e234eb4a928074dfd868cead0219ab85354cc53d20e772753c6bb9169d3 filter=lfs diff=lfs merge=lfs -text
 
 
 
 
38
  vae/ae.sft filter=lfs diff=lfs merge=lfs -text
39
  cache-huggingface/hub/models--QuanSun--EVA-CLIP/blobs/84c3a17a228c567a155259b2245b0b59072bf7da510260a0a02ec54de6d50b05 filter=lfs diff=lfs merge=lfs -text
40
  cache-huggingface/hub/models--bert-base-uncased/blobs/68d45e234eb4a928074dfd868cead0219ab85354cc53d20e772753c6bb9169d3 filter=lfs diff=lfs merge=lfs -text
41
+ lipsync-models/checkpoints/audio_model.ts filter=lfs diff=lfs merge=lfs -text
42
+ lipsync-models/checkpoints/lip_model.ts filter=lfs diff=lfs merge=lfs -text
43
+ lipsync-models/resemble_enhance/model_repo/enhancer_stage2/ds/G/default/mp_rank_00_model_states.pt.eeb346828f85451d959135c8db98d0cd.partial filter=lfs diff=lfs merge=lfs -text
lipsync-models/GPT_SoVITS/AR/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/AR/data/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/AR/data/bucket_sampler.py ADDED
@@ -0,0 +1,162 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/bucketsampler.py
2
+ import itertools
3
+ import math
4
+ import random
5
+ from random import shuffle
6
+ from typing import Iterator
7
+ from typing import Optional
8
+ from typing import TypeVar
9
+
10
+ import torch
11
+ import torch.distributed as dist
12
+ from torch.utils.data import Dataset
13
+ from torch.utils.data import Sampler
14
+
15
+ __all__ = [
16
+ "DistributedBucketSampler",
17
+ ]
18
+
19
+ T_co = TypeVar("T_co", covariant=True)
20
+
21
+
22
+ class DistributedBucketSampler(Sampler[T_co]):
23
+ r"""
24
+ sort the dataset wrt. input length
25
+ divide samples into buckets
26
+ sort within buckets
27
+ divide buckets into batches
28
+ sort batches
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ dataset: Dataset,
34
+ num_replicas: Optional[int] = None,
35
+ rank: Optional[int] = None,
36
+ shuffle: bool = True,
37
+ seed: int = 0,
38
+ drop_last: bool = False,
39
+ batch_size: int = 32,
40
+ ) -> None:
41
+ if num_replicas is None:
42
+ if not dist.is_available():
43
+ raise RuntimeError("Requires distributed package to be available")
44
+ num_replicas = dist.get_world_size() if torch.cuda.is_available() else 1
45
+ if rank is None:
46
+ if not dist.is_available():
47
+ raise RuntimeError("Requires distributed package to be available")
48
+ rank = dist.get_rank() if torch.cuda.is_available() else 0
49
+ if torch.cuda.is_available():
50
+ torch.cuda.set_device(rank)
51
+ if rank >= num_replicas or rank < 0:
52
+ raise ValueError(
53
+ "Invalid rank {}, rank should be in the interval"
54
+ " [0, {}]".format(rank, num_replicas - 1)
55
+ )
56
+ self.dataset = dataset
57
+ self.num_replicas = num_replicas
58
+ self.rank = rank
59
+ self.epoch = 0
60
+ self.drop_last = drop_last
61
+ # If the dataset length is evenly divisible by # of replicas, then there
62
+ # is no need to drop any data, since the dataset will be split equally.
63
+ if (
64
+ self.drop_last and len(self.dataset) % self.num_replicas != 0
65
+ ): # type: ignore[arg-type]
66
+ # Split to nearest available length that is evenly divisible.
67
+ # This is to ensure each rank receives the same amount of data when
68
+ # using this Sampler.
69
+ self.num_samples = math.ceil(
70
+ (len(self.dataset) - self.num_replicas)
71
+ / self.num_replicas # type: ignore[arg-type]
72
+ )
73
+ else:
74
+ self.num_samples = math.ceil(
75
+ len(self.dataset) / self.num_replicas
76
+ ) # type: ignore[arg-type]
77
+ self.total_size = self.num_samples * self.num_replicas
78
+ self.shuffle = shuffle
79
+ self.seed = seed
80
+ self.batch_size = batch_size
81
+ self.id_with_length = self._get_sample_lengths()
82
+ self.id_buckets = self.make_buckets(bucket_width=2.0)
83
+
84
+ def _get_sample_lengths(self):
85
+ id_with_lengths = []
86
+ for i in range(len(self.dataset)):
87
+ id_with_lengths.append((i, self.dataset.get_sample_length(i)))
88
+ id_with_lengths.sort(key=lambda x: x[1])
89
+ return id_with_lengths
90
+
91
+ def make_buckets(self, bucket_width: float = 2.0):
92
+ buckets = []
93
+ cur = []
94
+ max_sec = bucket_width
95
+ for id, sec in self.id_with_length:
96
+ if sec < max_sec:
97
+ cur.append(id)
98
+ else:
99
+ buckets.append(cur)
100
+ cur = [id]
101
+ max_sec += bucket_width
102
+ if len(cur) > 0:
103
+ buckets.append(cur)
104
+ return buckets
105
+
106
+ def __iter__(self) -> Iterator[T_co]:
107
+ if self.shuffle:
108
+ # deterministically shuffle based on epoch and seed
109
+ g = torch.Generator()
110
+ g.manual_seed(self.seed + self.epoch)
111
+ random.seed(self.epoch + self.seed)
112
+ shuffled_bucket = []
113
+ for buc in self.id_buckets:
114
+ buc_copy = buc.copy()
115
+ shuffle(buc_copy)
116
+ shuffled_bucket.append(buc_copy)
117
+ grouped_batch_size = self.batch_size * self.num_replicas
118
+ shuffled_bucket = list(itertools.chain(*shuffled_bucket))
119
+ n_batch = int(math.ceil(len(shuffled_bucket) / grouped_batch_size))
120
+ batches = [
121
+ shuffled_bucket[b * grouped_batch_size : (b + 1) * grouped_batch_size]
122
+ for b in range(n_batch)
123
+ ]
124
+ shuffle(batches)
125
+ indices = list(itertools.chain(*batches))
126
+ else:
127
+ # type: ignore[arg-type]
128
+ indices = list(range(len(self.dataset)))
129
+
130
+ if not self.drop_last:
131
+ # add extra samples to make it evenly divisible
132
+ padding_size = self.total_size - len(indices)
133
+ if padding_size <= len(indices):
134
+ indices += indices[:padding_size]
135
+ else:
136
+ indices += (indices * math.ceil(padding_size / len(indices)))[
137
+ :padding_size
138
+ ]
139
+ else:
140
+ # remove tail of data to make it evenly divisible.
141
+ indices = indices[: self.total_size]
142
+ assert len(indices) == self.total_size
143
+
144
+ # subsample
145
+ indices = indices[self.rank : self.total_size : self.num_replicas]
146
+ assert len(indices) == self.num_samples
147
+
148
+ return iter(indices)
149
+
150
+ def __len__(self) -> int:
151
+ return self.num_samples
152
+
153
+ def set_epoch(self, epoch: int) -> None:
154
+ r"""
155
+ Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
156
+ use a different random ordering for each epoch. Otherwise, the next iteration of this
157
+ sampler will yield the same ordering.
158
+
159
+ Args:
160
+ epoch (int): Epoch number.
161
+ """
162
+ self.epoch = epoch
lipsync-models/GPT_SoVITS/AR/data/data_module.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/data_module.py
2
+ from pytorch_lightning import LightningDataModule
3
+ from AR.data.bucket_sampler import DistributedBucketSampler
4
+ from AR.data.dataset import Text2SemanticDataset
5
+ from torch.utils.data import DataLoader
6
+
7
+
8
+ class Text2SemanticDataModule(LightningDataModule):
9
+ def __init__(
10
+ self,
11
+ config,
12
+ train_semantic_path,
13
+ train_phoneme_path,
14
+ dev_semantic_path=None,
15
+ dev_phoneme_path=None,
16
+ ):
17
+ super().__init__()
18
+ self.config = config
19
+ self.train_semantic_path = train_semantic_path
20
+ self.train_phoneme_path = train_phoneme_path
21
+ self.dev_semantic_path = dev_semantic_path
22
+ self.dev_phoneme_path = dev_phoneme_path
23
+ self.num_workers = self.config["data"]["num_workers"]
24
+
25
+ def prepare_data(self):
26
+ pass
27
+
28
+ def setup(self, stage=None, output_logs=False):
29
+ self._train_dataset = Text2SemanticDataset(
30
+ phoneme_path=self.train_phoneme_path,
31
+ semantic_path=self.train_semantic_path,
32
+ max_sec=self.config["data"]["max_sec"],
33
+ pad_val=self.config["data"]["pad_val"],
34
+ )
35
+ self._dev_dataset = self._train_dataset
36
+ # self._dev_dataset = Text2SemanticDataset(
37
+ # phoneme_path=self.dev_phoneme_path,
38
+ # semantic_path=self.dev_semantic_path,
39
+ # max_sample=self.config['data']['max_eval_sample'],
40
+ # max_sec=self.config['data']['max_sec'],
41
+ # pad_val=self.config['data']['pad_val'])
42
+
43
+ def train_dataloader(self):
44
+ batch_size=self.config["train"]["batch_size"]//2 if self.config["train"].get("if_dpo",False)==True else self.config["train"]["batch_size"]
45
+ batch_size = max(min(batch_size,len(self._train_dataset)//4),1)#防止不保存
46
+ sampler = DistributedBucketSampler(self._train_dataset, batch_size=batch_size)
47
+ return DataLoader(
48
+ self._train_dataset,
49
+ batch_size=batch_size,
50
+ sampler=sampler,
51
+ collate_fn=self._train_dataset.collate,
52
+ num_workers=self.num_workers,
53
+ persistent_workers=True,
54
+ prefetch_factor=16,
55
+ )
56
+
57
+ def val_dataloader(self):
58
+ return DataLoader(
59
+ self._dev_dataset,
60
+ batch_size=1,
61
+ shuffle=False,
62
+ collate_fn=self._train_dataset.collate,
63
+ num_workers=max(self.num_workers, 12),
64
+ persistent_workers=True,
65
+ prefetch_factor=16,
66
+ )
67
+
68
+ # 这个会使用到嘛?
69
+ def test_dataloader(self):
70
+ return DataLoader(
71
+ self._dev_dataset,
72
+ batch_size=1,
73
+ shuffle=False,
74
+ collate_fn=self._train_dataset.collate,
75
+ )
lipsync-models/GPT_SoVITS/AR/data/dataset.py ADDED
@@ -0,0 +1,320 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/t2s_dataset.py
2
+ import pdb
3
+ import sys
4
+
5
+ # sys.path.append("/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert")
6
+ import traceback, os
7
+ from typing import Dict
8
+ from typing import List
9
+
10
+ import numpy as np
11
+ import pandas as pd
12
+ import torch, json
13
+ from torch.utils.data import DataLoader
14
+ from torch.utils.data import Dataset
15
+ from transformers import AutoTokenizer
16
+
17
+ from text import cleaned_text_to_sequence
18
+
19
+ # from config import exp_dir
20
+
21
+
22
+ def batch_sequences(sequences: List[np.array], axis: int = 0, pad_value: int = 0):
23
+ seq = sequences[0]
24
+ ndim = seq.ndim
25
+ if axis < 0:
26
+ axis += ndim
27
+ dtype = seq.dtype
28
+ pad_value = dtype.type(pad_value)
29
+ seq_lengths = [seq.shape[axis] for seq in sequences]
30
+ max_length = np.max(seq_lengths)
31
+
32
+ padded_sequences = []
33
+ for seq, length in zip(sequences, seq_lengths):
34
+ padding = (
35
+ [(0, 0)] * axis + [(0, max_length - length)] + [(0, 0)] * (ndim - axis - 1)
36
+ )
37
+ padded_seq = np.pad(seq, padding, mode="constant", constant_values=pad_value)
38
+ padded_sequences.append(padded_seq)
39
+ batch = np.stack(padded_sequences)
40
+ return batch
41
+
42
+
43
+ class Text2SemanticDataset(Dataset):
44
+ """dataset class for text tokens to semantic model training."""
45
+
46
+ def __init__(
47
+ self,
48
+ phoneme_path: str,
49
+ semantic_path: str,
50
+ max_sample: int = None,
51
+ max_sec: int = 100,
52
+ pad_val: int = 1024,
53
+ # min value of phoneme/sec
54
+ min_ps_ratio: int = 3,
55
+ # max value of phoneme/sec
56
+ max_ps_ratio: int = 25,
57
+ ) -> None:
58
+ super().__init__()
59
+
60
+ self.semantic_data = pd.read_csv(
61
+ semantic_path, delimiter="\t", encoding="utf-8"
62
+ )
63
+ # get dict
64
+ self.path2 = phoneme_path # "%s/2-name2text.txt"%exp_dir#phoneme_path
65
+ self.path3 = "%s/3-bert" % (
66
+ os.path.basename(phoneme_path)
67
+ ) # "%s/3-bert"%exp_dir#bert_dir
68
+ self.path6 = semantic_path # "%s/6-name2semantic.tsv"%exp_dir#semantic_path
69
+ assert os.path.exists(self.path2)
70
+ assert os.path.exists(self.path6)
71
+ self.phoneme_data = {}
72
+ with open(self.path2, "r", encoding="utf8") as f:
73
+ lines = f.read().strip("\n").split("\n")
74
+
75
+ for line in lines:
76
+ tmp = line.split("\t")
77
+ if len(tmp) != 4:
78
+ continue
79
+ self.phoneme_data[tmp[0]] = [tmp[1], tmp[2], tmp[3]]
80
+
81
+ # self.phoneme_data = np.load(phoneme_path, allow_pickle=True).item()
82
+ # pad for semantic tokens
83
+ self.PAD: int = pad_val
84
+ # self.hz = 25
85
+ # with open("/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert/configs/s2.json", "r") as f:data = f.read()
86
+ # data=json.loads(data)["model"]["semantic_frame_rate"]#50hz
87
+ # self.hz=int(data[:-2])#
88
+ self.hz = int(os.environ.get("hz", "25hz")[:-2])
89
+
90
+ # max seconds of semantic token
91
+ self.max_sec = max_sec
92
+ self.min_ps_ratio = min_ps_ratio
93
+ self.max_ps_ratio = max_ps_ratio
94
+
95
+ if max_sample is not None:
96
+ self.semantic_data = self.semantic_data[:max_sample]
97
+
98
+ # {idx: (semantic, phoneme)}
99
+ # semantic list, phoneme list
100
+ self.semantic_phoneme = []
101
+ self.item_names = []
102
+
103
+ self.inited = False
104
+
105
+ if not self.inited:
106
+ # 调用初始化函数
107
+ self.init_batch()
108
+ self.inited = True
109
+ del self.semantic_data
110
+ del self.phoneme_data
111
+ # self.tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext-large")
112
+ # self.tokenizer = AutoTokenizer.from_pretrained("/data/docker/liujing04/bert-vits2/Bert-VITS2-master20231106/bert/chinese-roberta-wwm-ext-large")
113
+
114
+ def init_batch(self):
115
+ semantic_data_len = len(self.semantic_data)
116
+ phoneme_data_len = len(self.phoneme_data.keys())
117
+ print("semantic_data_len:", semantic_data_len)
118
+ print("phoneme_data_len:", phoneme_data_len)
119
+ print(self.semantic_data)
120
+ idx = 0
121
+ num_not_in = 0
122
+ num_deleted_bigger = 0
123
+ num_deleted_ps = 0
124
+ for i in range(semantic_data_len):
125
+ # 先依次遍历
126
+ # get str
127
+ item_name = self.semantic_data.iloc[i,0]
128
+ # print(self.phoneme_data)
129
+ try:
130
+ phoneme, word2ph, text = self.phoneme_data[item_name]
131
+ except Exception:
132
+ traceback.print_exc()
133
+ # print(f"{item_name} not in self.phoneme_data !")
134
+ num_not_in += 1
135
+ continue
136
+
137
+ semantic_str = self.semantic_data.iloc[i,1]
138
+ # get token list
139
+ semantic_ids = [int(idx) for idx in semantic_str.split(" ")]
140
+ # (T), 是否需要变成 (1, T) -> 不需要,因为需要求 len
141
+ # 过滤掉太长的样本
142
+ if (
143
+ len(semantic_ids) > self.max_sec * self.hz
144
+ ): #########1###根据token个数推测总时长过滤时长60s(config里)#40*25=1k
145
+ num_deleted_bigger += 1
146
+ continue
147
+ # (T, ), 这个速度不会很慢,所以可以在一开始就处理,无需在 __getitem__ 里面单个处理####
148
+ phoneme = phoneme.split(" ")
149
+
150
+ try:
151
+ phoneme_ids = cleaned_text_to_sequence(phoneme)
152
+ except:
153
+ traceback.print_exc()
154
+ # print(f"{item_name} not in self.phoneme_data !")
155
+ num_not_in += 1
156
+ continue
157
+ # if len(phoneme_ids) >400:###########2:改为恒定限制为semantic/2.5就行
158
+ if (
159
+ len(phoneme_ids) > self.max_sec * self.hz / 2.5
160
+ ): ###########2:改为恒定限制为semantic/2.5就行
161
+ num_deleted_ps += 1
162
+ continue
163
+ # if len(semantic_ids) > 1000:###########3
164
+ # num_deleted_bigger += 1
165
+ # continue
166
+
167
+ ps_ratio = len(phoneme_ids) / (len(semantic_ids) / self.hz)
168
+
169
+ if (
170
+ ps_ratio > self.max_ps_ratio or ps_ratio < self.min_ps_ratio
171
+ ): ##########4#3~25#每秒多少个phone
172
+ num_deleted_ps += 1
173
+ # print(item_name)
174
+ continue
175
+
176
+ self.semantic_phoneme.append((semantic_ids, phoneme_ids))
177
+ idx += 1
178
+ self.item_names.append(item_name)
179
+
180
+ min_num = 100 # 20直接不补#30补了也不存ckpt
181
+ leng = len(self.semantic_phoneme)
182
+ if leng < min_num:
183
+ tmp1 = self.semantic_phoneme
184
+ tmp2 = self.item_names
185
+ self.semantic_phoneme = []
186
+ self.item_names = []
187
+ for _ in range(max(2, int(min_num / leng))):
188
+ self.semantic_phoneme += tmp1
189
+ self.item_names += tmp2
190
+ if num_not_in > 0:
191
+ print(f"there are {num_not_in} semantic datas not in phoneme datas")
192
+ if num_deleted_bigger > 0:
193
+ print(
194
+ f"deleted {num_deleted_bigger} audios who's duration are bigger than {self.max_sec} seconds"
195
+ )
196
+ if num_deleted_ps > 0:
197
+ # 4702 for LibriTTS, LirbriTTS 是标注数据, 是否需要筛?=> 需要,有值为 100 的极端值
198
+ print(
199
+ f"deleted {num_deleted_ps} audios who's phoneme/sec are bigger than {self.max_ps_ratio} or smaller than {self.min_ps_ratio}"
200
+ )
201
+ """
202
+ there are 31 semantic datas not in phoneme datas
203
+ deleted 34 audios who's duration are bigger than 54 seconds
204
+ deleted 3190 audios who's phoneme/sec are bigger than 25 or smaller than 3
205
+ dataset.__len__(): 366463
206
+
207
+ """
208
+ # 345410 for LibriTTS
209
+ print("dataset.__len__():", self.__len__())
210
+
211
+ def __get_item_names__(self) -> List[str]:
212
+ return self.item_names
213
+
214
+ def __len__(self) -> int:
215
+ return len(self.semantic_phoneme)
216
+
217
+ def __getitem__(self, idx: int) -> Dict:
218
+ semantic_ids, phoneme_ids = self.semantic_phoneme[idx]
219
+ item_name = self.item_names[idx]
220
+ phoneme_ids_len = len(phoneme_ids)
221
+ # semantic tokens target
222
+ semantic_ids_len = len(semantic_ids)
223
+
224
+ flag = 0
225
+ path_bert = "%s/%s.pt" % (self.path3, item_name)
226
+ if os.path.exists(path_bert) == True:
227
+ bert_feature = torch.load(path_bert, map_location="cpu")
228
+ else:
229
+ flag = 1
230
+ if flag == 1:
231
+ # bert_feature=torch.zeros_like(phoneme_ids,dtype=torch.float32)
232
+ bert_feature = None
233
+ else:
234
+ assert bert_feature.shape[-1] == len(phoneme_ids)
235
+ return {
236
+ "idx": idx,
237
+ "phoneme_ids": phoneme_ids,
238
+ "phoneme_ids_len": phoneme_ids_len,
239
+ "semantic_ids": semantic_ids,
240
+ "semantic_ids_len": semantic_ids_len,
241
+ "bert_feature": bert_feature,
242
+ }
243
+
244
+ def get_sample_length(self, idx: int):
245
+ semantic_ids = self.semantic_phoneme[idx][0]
246
+ sec = 1.0 * len(semantic_ids) / self.hz
247
+ return sec
248
+
249
+ def collate(self, examples: List[Dict]) -> Dict:
250
+ sample_index: List[int] = []
251
+ phoneme_ids: List[torch.Tensor] = []
252
+ phoneme_ids_lens: List[int] = []
253
+ semantic_ids: List[torch.Tensor] = []
254
+ semantic_ids_lens: List[int] = []
255
+ # return
256
+
257
+ for item in examples:
258
+ sample_index.append(item["idx"])
259
+ phoneme_ids.append(np.array(item["phoneme_ids"], dtype=np.int64))
260
+ semantic_ids.append(np.array(item["semantic_ids"], dtype=np.int64))
261
+ phoneme_ids_lens.append(item["phoneme_ids_len"])
262
+ semantic_ids_lens.append(item["semantic_ids_len"])
263
+
264
+ # pad 0
265
+ phoneme_ids = batch_sequences(phoneme_ids)
266
+ semantic_ids = batch_sequences(semantic_ids, pad_value=self.PAD)
267
+
268
+ # # convert each batch to torch.tensor
269
+ phoneme_ids = torch.tensor(phoneme_ids)
270
+ semantic_ids = torch.tensor(semantic_ids)
271
+ phoneme_ids_lens = torch.tensor(phoneme_ids_lens)
272
+ semantic_ids_lens = torch.tensor(semantic_ids_lens)
273
+ bert_padded = torch.FloatTensor(len(examples), 1024, max(phoneme_ids_lens))
274
+ bert_padded.zero_()
275
+
276
+ for idx, item in enumerate(examples):
277
+ bert = item["bert_feature"]
278
+ if bert != None:
279
+ bert_padded[idx, :, : bert.shape[-1]] = bert
280
+
281
+ return {
282
+ # List[int]
283
+ "ids": sample_index,
284
+ # torch.Tensor (B, max_phoneme_length)
285
+ "phoneme_ids": phoneme_ids,
286
+ # torch.Tensor (B)
287
+ "phoneme_ids_len": phoneme_ids_lens,
288
+ # torch.Tensor (B, max_semantic_ids_length)
289
+ "semantic_ids": semantic_ids,
290
+ # torch.Tensor (B)
291
+ "semantic_ids_len": semantic_ids_lens,
292
+ # torch.Tensor (B, 1024, max_phoneme_length)
293
+ "bert_feature": bert_padded,
294
+ }
295
+
296
+
297
+ if __name__ == "__main__":
298
+ root_dir = "/data/docker/liujing04/gpt-vits/prepare/dump_mix/"
299
+ dataset = Text2SemanticDataset(
300
+ phoneme_path=root_dir + "phoneme_train.npy",
301
+ semantic_path=root_dir + "semantic_train.tsv",
302
+ )
303
+
304
+ batch_size = 12
305
+ dataloader = DataLoader(
306
+ dataset, batch_size=batch_size, collate_fn=dataset.collate, shuffle=False
307
+ )
308
+ for i, batch in enumerate(dataloader):
309
+ if i % 1000 == 0:
310
+ print(i)
311
+ # if i == 0:
312
+ # print('batch["ids"]:', batch["ids"])
313
+ # print('batch["phoneme_ids"]:', batch["phoneme_ids"],
314
+ # batch["phoneme_ids"].shape)
315
+ # print('batch["phoneme_ids_len"]:', batch["phoneme_ids_len"],
316
+ # batch["phoneme_ids_len"].shape)
317
+ # print('batch["semantic_ids"]:', batch["semantic_ids"],
318
+ # batch["semantic_ids"].shape)
319
+ # print('batch["semantic_ids_len"]:', batch["semantic_ids_len"],
320
+ # batch["semantic_ids_len"].shape)
lipsync-models/GPT_SoVITS/AR/models/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/AR/models/t2s_lightning_module.py ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_lightning_module.py
2
+ import os, sys
3
+
4
+ now_dir = os.getcwd()
5
+ sys.path.append(now_dir)
6
+ from typing import Dict
7
+
8
+ import torch
9
+ from pytorch_lightning import LightningModule
10
+ from AR.models.t2s_model import Text2SemanticDecoder
11
+ from AR.modules.lr_schedulers import WarmupCosineLRSchedule
12
+ from AR.modules.optim import ScaledAdam
13
+
14
+ class Text2SemanticLightningModule(LightningModule):
15
+ def __init__(self, config, output_dir, is_train=True):
16
+ super().__init__()
17
+ self.config = config
18
+ self.top_k = 3
19
+ self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)
20
+ pretrained_s1 = config.get("pretrained_s1")
21
+ if pretrained_s1 and is_train:
22
+ # print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["state_dict"]))
23
+ print(
24
+ self.load_state_dict(
25
+ torch.load(pretrained_s1, map_location="cpu")["weight"]
26
+ )
27
+ )
28
+ if is_train:
29
+ self.automatic_optimization = False
30
+ self.save_hyperparameters()
31
+ self.eval_dir = output_dir / "eval"
32
+ self.eval_dir.mkdir(parents=True, exist_ok=True)
33
+
34
+ def training_step(self, batch: Dict, batch_idx: int):
35
+ opt = self.optimizers()
36
+ scheduler = self.lr_schedulers()
37
+ forward=self.model.forward if self.config["train"].get("if_dpo",False)==True else self.model.forward_old
38
+ loss, acc = forward(
39
+ batch["phoneme_ids"],
40
+ batch["phoneme_ids_len"],
41
+ batch["semantic_ids"],
42
+ batch["semantic_ids_len"],
43
+ batch["bert_feature"],
44
+ )
45
+ self.manual_backward(loss)
46
+ if batch_idx > 0 and batch_idx % 4 == 0:
47
+ opt.step()
48
+ opt.zero_grad()
49
+ scheduler.step()
50
+
51
+ self.log(
52
+ "total_loss",
53
+ loss,
54
+ on_step=True,
55
+ on_epoch=True,
56
+ prog_bar=True,
57
+ sync_dist=True,
58
+ )
59
+ self.log(
60
+ "lr",
61
+ scheduler.get_last_lr()[0],
62
+ on_epoch=True,
63
+ prog_bar=True,
64
+ sync_dist=True,
65
+ )
66
+ self.log(
67
+ f"top_{self.top_k}_acc",
68
+ acc,
69
+ on_step=True,
70
+ on_epoch=True,
71
+ prog_bar=True,
72
+ sync_dist=True,
73
+ )
74
+
75
+ def validation_step(self, batch: Dict, batch_idx: int):
76
+ return
77
+
78
+ # # get loss
79
+ # loss, acc = self.model.forward(
80
+ # batch['phoneme_ids'], batch['phoneme_ids_len'],
81
+ # batch['semantic_ids'], batch['semantic_ids_len'],
82
+ # batch['bert_feature']
83
+ # )
84
+ #
85
+ # self.log(
86
+ # "val_total_loss",
87
+ # loss,
88
+ # on_step=True,
89
+ # on_epoch=True,
90
+ # prog_bar=True,
91
+ # sync_dist=True)
92
+ # self.log(
93
+ # f"val_top_{self.top_k}_acc",
94
+ # acc,
95
+ # on_step=True,
96
+ # on_epoch=True,
97
+ # prog_bar=True,
98
+ # sync_dist=True)
99
+ #
100
+ # # get infer output
101
+ # semantic_len = batch['semantic_ids'].size(1)
102
+ # prompt_len = min(int(semantic_len * 0.5), 150)
103
+ # prompt = batch['semantic_ids'][:, :prompt_len]
104
+ # pred_semantic = self.model.infer(batch['phoneme_ids'],
105
+ # batch['phoneme_ids_len'], prompt,
106
+ # batch['bert_feature']
107
+ # )
108
+ # save_name = f'semantic_toks_{batch_idx}.pt'
109
+ # save_path = os.path.join(self.eval_dir, save_name)
110
+ # torch.save(pred_semantic.detach().cpu(), save_path)
111
+
112
+ def configure_optimizers(self):
113
+ model_parameters = self.model.parameters()
114
+ parameters_names = []
115
+ parameters_names.append(
116
+ [name_param_pair[0] for name_param_pair in self.model.named_parameters()]
117
+ )
118
+ lm_opt = ScaledAdam(
119
+ model_parameters,
120
+ lr=0.01,
121
+ betas=(0.9, 0.95),
122
+ clipping_scale=2.0,
123
+ parameters_names=parameters_names,
124
+ show_dominant_parameters=False,
125
+ clipping_update_period=1000,
126
+ )
127
+
128
+ return {
129
+ "optimizer": lm_opt,
130
+ "lr_scheduler": {
131
+ "scheduler": WarmupCosineLRSchedule(
132
+ lm_opt,
133
+ init_lr=self.config["optimizer"]["lr_init"],
134
+ peak_lr=self.config["optimizer"]["lr"],
135
+ end_lr=self.config["optimizer"]["lr_end"],
136
+ warmup_steps=self.config["optimizer"]["warmup_steps"],
137
+ total_steps=self.config["optimizer"]["decay_steps"],
138
+ )
139
+ },
140
+ }
lipsync-models/GPT_SoVITS/AR/models/t2s_model.py ADDED
@@ -0,0 +1,361 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py
2
+ # reference: https://github.com/lifeiteng/vall-e
3
+
4
+ import torch
5
+ from AR.models.utils import (
6
+ dpo_loss,
7
+ get_batch_logps,
8
+ logits_to_probs,
9
+ make_pad_mask,
10
+ make_reject_y,
11
+ multinomial_sample_one_no_sync,
12
+ sample,
13
+ topk_sampling,
14
+ )
15
+ from AR.modules.embedding import SinePositionalEmbedding, TokenEmbedding
16
+ from AR.modules.transformer import (
17
+ LayerNorm,
18
+ TransformerEncoder,
19
+ TransformerEncoderLayer,
20
+ )
21
+ from typing import List, Optional
22
+ from torch import nn
23
+ from torch.nn import functional as F
24
+ from tqdm import tqdm
25
+
26
+ default_config = {
27
+ "embedding_dim": 512,
28
+ "hidden_dim": 512,
29
+ "num_head": 8,
30
+ "num_layers": 12,
31
+ "num_codebook": 8,
32
+ "p_dropout": 0.0,
33
+ "vocab_size": 1024 + 1,
34
+ "phoneme_vocab_size": 512,
35
+ "EOS": 1024,
36
+ }
37
+
38
+
39
+ @torch.jit.script
40
+ class T2SMLP:
41
+ def __init__(self, w1, b1, w2, b2):
42
+ self.w1 = w1
43
+ self.b1 = b1
44
+ self.w2 = w2
45
+ self.b2 = b2
46
+
47
+ def forward(self, x):
48
+ x = F.relu(F.linear(x, self.w1, self.b1))
49
+ x = F.linear(x, self.w2, self.b2)
50
+ return x
51
+
52
+
53
+ @torch.jit.script
54
+ class T2SBlock:
55
+ def __init__(
56
+ self,
57
+ num_heads: int,
58
+ hidden_dim: int,
59
+ mlp: T2SMLP,
60
+ qkv_w,
61
+ qkv_b,
62
+ out_w,
63
+ out_b,
64
+ norm_w1,
65
+ norm_b1,
66
+ norm_eps1: float,
67
+ norm_w2,
68
+ norm_b2,
69
+ norm_eps2: float,
70
+ ):
71
+ self.num_heads = num_heads
72
+ self.mlp = mlp
73
+ self.hidden_dim: int = hidden_dim
74
+ self.qkv_w = qkv_w
75
+ self.qkv_b = qkv_b
76
+ self.out_w = out_w
77
+ self.out_b = out_b
78
+ self.norm_w1 = norm_w1
79
+ self.norm_b1 = norm_b1
80
+ self.norm_eps1 = norm_eps1
81
+ self.norm_w2 = norm_w2
82
+ self.norm_b2 = norm_b2
83
+ self.norm_eps2 = norm_eps2
84
+
85
+ def process_prompt(self, x, attn_mask: torch.Tensor):
86
+ q, k, v = F.linear(x, self.qkv_w, self.qkv_b).chunk(3, dim=-1)
87
+
88
+ batch_size = q.shape[0]
89
+ q_len = q.shape[1]
90
+ kv_len = k.shape[1]
91
+
92
+ k_cache = k
93
+ v_cache = v
94
+
95
+ q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2)
96
+ k = k_cache.view(batch_size, kv_len,
97
+ self.num_heads, -1).transpose(1, 2)
98
+ v = v_cache.view(batch_size, kv_len,
99
+ self.num_heads, -1).transpose(1, 2)
100
+
101
+ attn = F.scaled_dot_product_attention(q, k, v, ~attn_mask)
102
+
103
+ attn = attn.permute(2, 0, 1, 3).reshape(
104
+ batch_size, -1, self.hidden_dim)
105
+ attn = F.linear(attn, self.out_w, self.out_b)
106
+
107
+ x = F.layer_norm(
108
+ x + attn, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1
109
+ )
110
+ x = F.layer_norm(
111
+ x + self.mlp.forward(x),
112
+ [self.hidden_dim],
113
+ self.norm_w2,
114
+ self.norm_b2,
115
+ self.norm_eps2,
116
+ )
117
+ return x, k_cache, v_cache
118
+
119
+ def decode_next_token(self, x, k_cache, v_cache):
120
+ q, k, v = F.linear(x, self.qkv_w, self.qkv_b).chunk(3, dim=-1)
121
+
122
+ k_cache = torch.cat([k_cache, k], dim=1)
123
+ v_cache = torch.cat([v_cache, v], dim=1)
124
+ kv_len = k_cache.shape[1]
125
+
126
+ batch_size = q.shape[0]
127
+ q_len = q.shape[1]
128
+
129
+ q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2)
130
+ k = k_cache.view(batch_size, kv_len,
131
+ self.num_heads, -1).transpose(1, 2)
132
+ v = v_cache.view(batch_size, kv_len,
133
+ self.num_heads, -1).transpose(1, 2)
134
+
135
+ attn = F.scaled_dot_product_attention(q, k, v)
136
+
137
+ attn = attn.permute(2, 0, 1, 3).reshape(
138
+ batch_size, -1, self.hidden_dim)
139
+ attn = F.linear(attn, self.out_w, self.out_b)
140
+
141
+ x = F.layer_norm(
142
+ x + attn, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1
143
+ )
144
+ x = F.layer_norm(
145
+ x + self.mlp.forward(x),
146
+ [self.hidden_dim],
147
+ self.norm_w2,
148
+ self.norm_b2,
149
+ self.norm_eps2,
150
+ )
151
+ return x, k_cache, v_cache
152
+
153
+
154
+ @torch.jit.script
155
+ class T2STransformer:
156
+ def __init__(self, num_blocks: int, blocks: List[T2SBlock]):
157
+ self.num_blocks: int = num_blocks
158
+ self.blocks = blocks
159
+
160
+ def process_prompt(self, x, attn_mask: torch.Tensor):
161
+ k_cache: List[torch.Tensor] = []
162
+ v_cache: List[torch.Tensor] = []
163
+ for i in range(self.num_blocks):
164
+ x, k_cache_, v_cache_ = self.blocks[i].process_prompt(x, attn_mask)
165
+ k_cache.append(k_cache_)
166
+ v_cache.append(v_cache_)
167
+ return x, k_cache, v_cache
168
+
169
+ def decode_next_token(
170
+ self, x, k_cache: List[torch.Tensor], v_cache: List[torch.Tensor]
171
+ ):
172
+ for i in range(self.num_blocks):
173
+ x, k_cache[i], v_cache[i] = self.blocks[i].decode_next_token(
174
+ x, k_cache[i], v_cache[i]
175
+ )
176
+ return x, k_cache, v_cache
177
+
178
+
179
+ class Text2SemanticDecoder(nn.Module):
180
+ def __init__(self, config, norm_first=False, top_k=3):
181
+ super(Text2SemanticDecoder, self).__init__()
182
+ self.model_dim = config["model"]["hidden_dim"]
183
+ self.embedding_dim = config["model"]["embedding_dim"]
184
+ self.num_head = config["model"]["head"]
185
+ self.num_layers = config["model"]["n_layer"]
186
+ self.norm_first = norm_first
187
+ self.vocab_size = config["model"]["vocab_size"]
188
+ self.phoneme_vocab_size = config["model"]["phoneme_vocab_size"]
189
+ self.p_dropout = float(config["model"]["dropout"])
190
+ self.EOS = config["model"]["EOS"]
191
+ self.norm_first = norm_first
192
+ assert self.EOS == self.vocab_size - 1
193
+ # should be same as num of kmeans bin
194
+ # assert self.EOS == 1024
195
+ self.bert_proj = nn.Linear(1024, self.embedding_dim)
196
+ self.ar_text_embedding = TokenEmbedding(
197
+ self.embedding_dim, self.phoneme_vocab_size, self.p_dropout
198
+ )
199
+ self.ar_text_position = SinePositionalEmbedding(
200
+ self.embedding_dim, dropout=0.1, scale=False, alpha=True
201
+ )
202
+ self.ar_audio_embedding = TokenEmbedding(
203
+ self.embedding_dim, self.vocab_size, self.p_dropout
204
+ )
205
+ self.ar_audio_position = SinePositionalEmbedding(
206
+ self.embedding_dim, dropout=0.1, scale=False, alpha=True
207
+ )
208
+
209
+ self.h = TransformerEncoder(
210
+ TransformerEncoderLayer(
211
+ d_model=self.model_dim,
212
+ nhead=self.num_head,
213
+ dim_feedforward=self.model_dim * 4,
214
+ dropout=0.1,
215
+ batch_first=True,
216
+ norm_first=norm_first,
217
+ ),
218
+ num_layers=self.num_layers,
219
+ norm=LayerNorm(self.model_dim) if norm_first else None,
220
+ )
221
+
222
+ self.ar_predict_layer = nn.Linear(
223
+ self.model_dim, self.vocab_size, bias=False)
224
+
225
+ blocks = []
226
+
227
+ for i in range(self.num_layers):
228
+ layer = self.h.layers[i]
229
+ t2smlp = T2SMLP(
230
+ layer.linear1.weight,
231
+ layer.linear1.bias,
232
+ layer.linear2.weight,
233
+ layer.linear2.bias,
234
+ )
235
+ # (layer.self_attn.in_proj_weight, layer.self_attn.in_proj_bias)
236
+ block = T2SBlock(
237
+ self.num_head,
238
+ self.model_dim,
239
+ t2smlp,
240
+ layer.self_attn.in_proj_weight,
241
+ layer.self_attn.in_proj_bias,
242
+ layer.self_attn.out_proj.weight,
243
+ layer.self_attn.out_proj.bias,
244
+ layer.norm1.weight,
245
+ layer.norm1.bias,
246
+ layer.norm1.eps,
247
+ layer.norm2.weight,
248
+ layer.norm2.bias,
249
+ layer.norm2.eps,
250
+ )
251
+
252
+ blocks.append(block)
253
+
254
+ self.t2s_transformer = T2STransformer(self.num_layers, blocks)
255
+
256
+ def forward(
257
+ self,
258
+ seq: torch.LongTensor, # 全部文本token
259
+ prompts: torch.LongTensor, # 参考音频token
260
+ bert_feature,
261
+ top_k: int = 5,
262
+ top_p: int = 1,
263
+ early_stop_num: int = 2700,
264
+ temperature: float = 1.0,
265
+ ):
266
+ x = self.ar_text_embedding(seq)
267
+ x = x + self.bert_proj(bert_feature.transpose(1, 2))
268
+ x = self.ar_text_position(x)
269
+
270
+ # AR Decoder
271
+ y = prompts
272
+
273
+ x_len = x.shape[1]
274
+ x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
275
+ stop = False
276
+ # print(1111111,self.num_layers)
277
+
278
+ ################### first step ##########################
279
+ if y is not None:
280
+ y_emb = self.ar_audio_embedding(y)
281
+ y_len = y_emb.shape[1]
282
+ prefix_len = y.shape[1]
283
+ y_pos = self.ar_audio_position(y_emb)
284
+ xy_pos = torch.concat([x, y_pos], dim=1)
285
+ ref_free = False
286
+ else:
287
+ y_emb = None
288
+ y_len = 0
289
+ prefix_len = 0
290
+ y_pos = None
291
+ xy_pos = x
292
+ y = torch.zeros(x.shape[0], 0, dtype=torch.int, device=x.device)
293
+ ref_free = True
294
+
295
+ x_attn_mask_pad = F.pad(
296
+ x_attn_mask,
297
+ (0, y_len), # xx的纯0扩展到xx纯0+xy纯1,(x,x+y)
298
+ value=1.0,
299
+ )
300
+ y_attn_mask = F.pad( # yy的右上1扩展到左边xy的0,(y,x+y)
301
+ torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
302
+ (x_len, 0),
303
+ value=0.0,
304
+ )
305
+
306
+ xy_attn_mask = torch.concat(
307
+ [x_attn_mask_pad, y_attn_mask], dim=0).to(x.device)
308
+ xy_dec, k_cache, v_cache = self.t2s_transformer.process_prompt(
309
+ xy_pos, xy_attn_mask
310
+ )
311
+
312
+ final = 0
313
+
314
+ for idx in range(0, 1500):
315
+ if idx > 0:
316
+ xy_dec, k_cache, v_cache = self.t2s_transformer.decode_next_token(
317
+ xy_pos, k_cache, v_cache
318
+ )
319
+
320
+ logits = self.ar_predict_layer(xy_dec[:, -1])
321
+
322
+ if idx == 0:
323
+ logits = logits[:, :-1]
324
+
325
+ samples = sample(
326
+ logits[0],
327
+ y,
328
+ top_k=top_k,
329
+ top_p=top_p,
330
+ repetition_penalty=1.35,
331
+ temperature=temperature,
332
+ )[0].unsqueeze(0)
333
+
334
+ y = torch.concat([y, samples], dim=1)
335
+
336
+ if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:
337
+ print("use early stop num:", early_stop_num)
338
+ stop = True
339
+
340
+ if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
341
+ stop = True
342
+
343
+ if stop:
344
+ # if y.shape[1] == 0:
345
+ # y = torch.concat([y, torch.zeros_like(samples)], dim=1)
346
+ # print("bad zero prediction")
347
+ # print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]")
348
+ final = idx
349
+ break
350
+
351
+ ####################### update next step ###################################
352
+ y_emb = self.ar_audio_embedding(y[:, -1:])
353
+ xy_pos = (
354
+ y_emb * self.ar_audio_position.x_scale
355
+ + self.ar_audio_position.alpha
356
+ * self.ar_audio_position.pe[:, prompts.shape[1] + idx]
357
+ )
358
+
359
+ if ref_free:
360
+ return y[:, :-1], 0
361
+ return y[:, :-1], final - 1
lipsync-models/GPT_SoVITS/AR/models/utils.py ADDED
@@ -0,0 +1,243 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/utils.py\
2
+ from typing import Optional, Tuple
3
+ import torch
4
+ import torch.nn.functional as F
5
+ from typing import Tuple
6
+
7
+
8
+ def sequence_mask(length, max_length=None):
9
+ if max_length is None:
10
+ max_length = length.max()
11
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
12
+ return x.unsqueeze(0) < length.unsqueeze(1)
13
+
14
+
15
+ def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
16
+ """
17
+ Args:
18
+ lengths:
19
+ A 1-D tensor containing sentence lengths.
20
+ max_len:
21
+ The length of masks.
22
+ Returns:
23
+ Return a 2-D bool tensor, where masked positions
24
+ are filled with `True` and non-masked positions are
25
+ filled with `False`.
26
+
27
+ #>>> lengths = torch.tensor([1, 3, 2, 5])
28
+ #>>> make_pad_mask(lengths)
29
+ tensor([[False, True, True, True, True],
30
+ [False, False, False, True, True],
31
+ [False, False, True, True, True],
32
+ [False, False, False, False, False]])
33
+ """
34
+ assert lengths.ndim == 1, lengths.ndim
35
+ max_len = max(max_len, lengths.max())
36
+ n = lengths.size(0)
37
+ seq_range = torch.arange(0, max_len, device=lengths.device)
38
+ expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
39
+
40
+ return expaned_lengths >= lengths.unsqueeze(-1)
41
+
42
+
43
+ # https://github.com/microsoft/unilm/blob/master/xtune/src/transformers/modeling_utils.py
44
+ def top_k_top_p_filtering(
45
+ logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
46
+ ):
47
+ """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
48
+ Args:
49
+ logits: logits distribution shape (batch size, vocabulary size)
50
+ if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
51
+ if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
52
+ Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
53
+ Make sure we keep at least min_tokens_to_keep per batch example in the output
54
+ From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
55
+ """
56
+ if top_k > 0:
57
+ top_k = min(max(top_k, min_tokens_to_keep),
58
+ logits.size(-1)) # Safety check
59
+ # Remove all tokens with a probability less than the last token of the top-k
60
+ indices_to_remove = logits < torch.topk(logits, top_k)[
61
+ 0][..., -1, None]
62
+ logits[indices_to_remove] = filter_value
63
+
64
+ if top_p < 1.0:
65
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
66
+ cumulative_probs = torch.cumsum(
67
+ F.softmax(sorted_logits, dim=-1), dim=-1)
68
+
69
+ # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
70
+ sorted_indices_to_remove = cumulative_probs > top_p
71
+ if min_tokens_to_keep > 1:
72
+ # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
73
+ sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
74
+ # Shift the indices to the right to keep also the first token above the threshold
75
+ sorted_indices_to_remove[...,
76
+ 1:] = sorted_indices_to_remove[..., :-1].clone()
77
+ sorted_indices_to_remove[..., 0] = 0
78
+
79
+ # scatter sorted tensors to original indexing
80
+ indices_to_remove = sorted_indices_to_remove.scatter(
81
+ 1, sorted_indices, sorted_indices_to_remove
82
+ )
83
+ logits[indices_to_remove] = filter_value
84
+ return logits
85
+
86
+
87
+ def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
88
+ # temperature: (`optional`) float
89
+ # The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
90
+ # top_k: (`optional`) int
91
+ # The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
92
+ # top_p: (`optional`) float
93
+ # The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
94
+
95
+ # Temperature (higher temperature => more likely to sample low probability tokens)
96
+ if temperature != 1.0:
97
+ logits = logits / temperature
98
+ # Top-p/top-k filtering
99
+ logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
100
+ # Sample
101
+ token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
102
+ return token
103
+
104
+
105
+ def multinomial_sample_one_no_sync(
106
+ probs_sort,
107
+ ): # Does multinomial sampling without a cuda synchronization
108
+ q = torch.empty_like(probs_sort).exponential_(1.0)
109
+ return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
110
+
111
+
112
+ def logits_to_probs(
113
+ logits,
114
+ previous_tokens: Optional[torch.Tensor] = None,
115
+ temperature: float = 1.0,
116
+ top_k: Optional[int] = None,
117
+ top_p: Optional[int] = None,
118
+ repetition_penalty: float = 1.0,
119
+ ):
120
+ if previous_tokens is not None:
121
+ previous_tokens = previous_tokens.squeeze()
122
+ # print(logits.shape,previous_tokens.shape)
123
+ # pdb.set_trace()
124
+ if previous_tokens is not None and repetition_penalty != 1.0:
125
+ previous_tokens = previous_tokens.long()
126
+ score = torch.gather(logits, dim=0, index=previous_tokens)
127
+ score = torch.where(
128
+ score < 0, score * repetition_penalty, score / repetition_penalty
129
+ )
130
+ logits.scatter_(dim=0, index=previous_tokens, src=score)
131
+
132
+ if top_p is not None and top_p < 1.0:
133
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
134
+ cum_probs = torch.cumsum(
135
+ torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1
136
+ )
137
+ sorted_indices_to_remove = cum_probs > top_p
138
+ sorted_indices_to_remove[0] = False # keep at least one option
139
+ indices_to_remove = sorted_indices_to_remove.scatter(
140
+ dim=0, index=sorted_indices, src=sorted_indices_to_remove
141
+ )
142
+ logits = logits.masked_fill(indices_to_remove, -float("Inf"))
143
+
144
+ logits = logits / max(temperature, 1e-5)
145
+
146
+ if top_k is not None:
147
+ v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
148
+ pivot = v.select(-1, -1).unsqueeze(-1)
149
+ logits = torch.where(logits < pivot, -float("Inf"), logits)
150
+
151
+ probs = torch.nn.functional.softmax(logits, dim=-1)
152
+ return probs
153
+
154
+
155
+ def sample(
156
+ logits,
157
+ previous_tokens: Optional[torch.Tensor] = None,
158
+ temperature: float = 1.0,
159
+ top_k: Optional[int] = None,
160
+ top_p: Optional[int] = None,
161
+ repetition_penalty: float = 1.0,
162
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
163
+ probs = logits_to_probs(
164
+ logits=logits, previous_tokens=previous_tokens, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty
165
+ )
166
+ idx_next = multinomial_sample_one_no_sync(probs)
167
+ return idx_next, probs
168
+
169
+
170
+ def dpo_loss(policy_chosen_logps: torch.FloatTensor,
171
+ policy_rejected_logps: torch.FloatTensor,
172
+ reference_chosen_logps: torch.FloatTensor,
173
+ reference_rejected_logps: torch.FloatTensor,
174
+ beta: float,
175
+ reference_free: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
176
+ pi_logratios = policy_chosen_logps - policy_rejected_logps
177
+ ref_logratios = reference_chosen_logps - reference_rejected_logps
178
+
179
+ if reference_free:
180
+ ref_logratios = 0
181
+
182
+ logits = pi_logratios - ref_logratios
183
+
184
+ losses = -F.logsigmoid(beta * logits)
185
+ chosen_rewards = beta * (policy_chosen_logps -
186
+ reference_chosen_logps).detach()
187
+ rejected_rewards = beta * \
188
+ (policy_rejected_logps - reference_rejected_logps).detach()
189
+
190
+ return losses.mean(), chosen_rewards, rejected_rewards
191
+
192
+
193
+ def get_batch_logps(logits_target: torch.FloatTensor, logits_reject: torch.FloatTensor, labels_target: torch.LongTensor, labels_reject: torch.LongTensor, average_log_prob: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
194
+
195
+ # dummy token; we'll ignore the losses on these tokens later
196
+
197
+ per_token_logps_target = torch.gather(
198
+ logits_target.log_softmax(-1), dim=2, index=labels_target.unsqueeze(2)).squeeze(2)
199
+ per_token_logps_reject = torch.gather(
200
+ logits_reject.log_softmax(-1), dim=2, index=labels_reject.unsqueeze(2)).squeeze(2)
201
+
202
+ return per_token_logps_target.sum(-1), per_token_logps_reject.sum(-1)
203
+
204
+
205
+ def make_reject_y(y_o, y_lens):
206
+ def repeat_P(y):
207
+ range_idx, _ = torch.randint(0, len(y), size=(2,)).sort()
208
+ pre = y[:range_idx[0]]
209
+ shf = y[range_idx[1]:]
210
+ range_text = y[range_idx[0]:range_idx[1]]
211
+ new_y = torch.cat([pre, range_text, range_text, shf])
212
+ return new_y
213
+
214
+ def lost_P(y):
215
+ range_idx, _ = torch.randint(0, len(y), size=(2,)).sort()
216
+ pre = y[:range_idx[0]]
217
+ shf = y[range_idx[1]:]
218
+ range_text = y[range_idx[0]:range_idx[1]]
219
+ new_y = torch.cat([pre, shf])
220
+ return new_y
221
+ bs = len(y_lens)
222
+ reject_y = []
223
+ reject_y_lens = []
224
+ for b in range(bs):
225
+ process_item_idx = torch.randint(0, 1, size=(1, ))[0]
226
+ if process_item_idx == 0:
227
+ new_y = repeat_P(y_o[b])
228
+ reject_y.append(new_y)
229
+ reject_y_lens.append(len(new_y))
230
+ elif process_item_idx == 1:
231
+ new_y = lost_P(y_o[b])
232
+ reject_y.append(new_y)
233
+ reject_y_lens.append(len(new_y))
234
+ max_length = max(reject_y_lens)
235
+ for b in range(bs):
236
+ pad_length = max_length - reject_y_lens[b]
237
+ reject_y[b] = torch.cat([reject_y[b], torch.zeros(
238
+ pad_length, dtype=y_o.dtype, device=y_o.device)], dim=0)
239
+
240
+ reject_y = torch.stack(reject_y, dim=0)
241
+ reject_y_lens = torch.tensor(reject_y_lens, device=y_lens.device)
242
+
243
+ return reject_y, reject_y_lens
lipsync-models/GPT_SoVITS/AR/modules/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/AR/modules/activation.py ADDED
@@ -0,0 +1,426 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py
2
+ from typing import Optional
3
+ from typing import Tuple
4
+ import torch
5
+ from torch import Tensor
6
+ from torch.nn import Linear
7
+ from torch.nn import Module
8
+ from torch.nn.init import constant_
9
+ from torch.nn.init import xavier_normal_
10
+ from torch.nn.init import xavier_uniform_
11
+ from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
12
+ from torch.nn.parameter import Parameter
13
+
14
+ from torch.nn import functional as F
15
+ from AR.modules.patched_mha_with_cache import multi_head_attention_forward_patched
16
+
17
+
18
+ class MultiheadAttention(Module):
19
+ r"""Allows the model to jointly attend to information
20
+ from different representation subspaces as described in the paper:
21
+ `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
22
+
23
+ Multi-Head Attention is defined as:
24
+
25
+ .. math::
26
+ \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
27
+
28
+ where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
29
+
30
+ ``forward()`` will use a special optimized implementation if all of the following
31
+ conditions are met:
32
+
33
+ - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
34
+ restriction will be loosened in the future.)
35
+ - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
36
+ - training is disabled (using ``.eval()``)
37
+ - dropout is 0
38
+ - ``add_bias_kv`` is ``False``
39
+ - ``add_zero_attn`` is ``False``
40
+ - ``batch_first`` is ``True`` and the input is batched
41
+ - ``kdim`` and ``vdim`` are equal to ``embed_dim``
42
+ - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
43
+ - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
44
+ nor ``attn_mask`` is passed
45
+
46
+ If the optimized implementation is in use, a
47
+ `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
48
+ ``query``/``key``/``value`` to represent padding more efficiently than using a
49
+ padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
50
+ will be returned, and an additional speedup proportional to the fraction of the input
51
+ that is padding can be expected.
52
+
53
+ Args:
54
+ embed_dim: Total dimension of the model.
55
+ num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
56
+ across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
57
+ dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
58
+ bias: If specified, adds bias to input / output projection layers. Default: ``True``.
59
+ add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
60
+ add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
61
+ Default: ``False``.
62
+ kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
63
+ vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
64
+ batch_first: If ``True``, then the input and output tensors are provided
65
+ as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
66
+
67
+ Examples::
68
+
69
+ >>> # xdoctest: +SKIP
70
+ >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
71
+ >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
72
+
73
+ """
74
+ __constants__ = ["batch_first"]
75
+ bias_k: Optional[torch.Tensor]
76
+ bias_v: Optional[torch.Tensor]
77
+
78
+ def __init__(
79
+ self,
80
+ embed_dim,
81
+ num_heads,
82
+ dropout=0.0,
83
+ bias=True,
84
+ add_bias_kv=False,
85
+ add_zero_attn=False,
86
+ kdim=None,
87
+ vdim=None,
88
+ batch_first=False,
89
+ linear1_cls=Linear,
90
+ linear2_cls=Linear,
91
+ device=None,
92
+ dtype=None,
93
+ ) -> None:
94
+ factory_kwargs = {"device": device, "dtype": dtype}
95
+ super(MultiheadAttention, self).__init__()
96
+ self.embed_dim = embed_dim
97
+ self.kdim = kdim if kdim is not None else embed_dim
98
+ self.vdim = vdim if vdim is not None else embed_dim
99
+ self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
100
+
101
+ self.num_heads = num_heads
102
+ self.dropout = dropout
103
+ self.batch_first = batch_first
104
+ self.head_dim = embed_dim // num_heads
105
+ assert (
106
+ self.head_dim * num_heads == self.embed_dim
107
+ ), "embed_dim must be divisible by num_heads"
108
+
109
+ if add_bias_kv:
110
+ self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
111
+ self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs))
112
+ else:
113
+ self.bias_k = self.bias_v = None
114
+
115
+ if linear1_cls == Linear:
116
+ if not self._qkv_same_embed_dim:
117
+ self.q_proj_weight = Parameter(
118
+ torch.empty((embed_dim, embed_dim), **factory_kwargs)
119
+ )
120
+ self.k_proj_weight = Parameter(
121
+ torch.empty((embed_dim, self.kdim), **factory_kwargs)
122
+ )
123
+ self.v_proj_weight = Parameter(
124
+ torch.empty((embed_dim, self.vdim), **factory_kwargs)
125
+ )
126
+ self.register_parameter("in_proj_weight", None)
127
+ else:
128
+ self.in_proj_weight = Parameter(
129
+ torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
130
+ )
131
+ self.register_parameter("q_proj_weight", None)
132
+ self.register_parameter("k_proj_weight", None)
133
+ self.register_parameter("v_proj_weight", None)
134
+
135
+ if bias:
136
+ self.in_proj_bias = Parameter(
137
+ torch.empty(3 * embed_dim, **factory_kwargs)
138
+ )
139
+ else:
140
+ self.register_parameter("in_proj_bias", None)
141
+ self.out_proj = NonDynamicallyQuantizableLinear(
142
+ embed_dim, embed_dim, bias=bias, **factory_kwargs
143
+ )
144
+
145
+ self._reset_parameters()
146
+ else:
147
+ if not self._qkv_same_embed_dim:
148
+ raise NotImplementedError
149
+ else:
150
+ self.in_proj_linear = linear1_cls(
151
+ embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
152
+ )
153
+ self.in_proj_weight = self.in_proj_linear.weight
154
+
155
+ self.register_parameter("q_proj_weight", None)
156
+ self.register_parameter("k_proj_weight", None)
157
+ self.register_parameter("v_proj_weight", None)
158
+
159
+ if bias:
160
+ self.in_proj_bias = self.in_proj_linear.bias
161
+ else:
162
+ self.register_parameter("in_proj_bias", None)
163
+
164
+ self.out_proj = linear2_cls(
165
+ embed_dim, embed_dim, bias=bias, **factory_kwargs
166
+ )
167
+
168
+ if self.bias_k is not None:
169
+ xavier_normal_(self.bias_k)
170
+ if self.bias_v is not None:
171
+ xavier_normal_(self.bias_v)
172
+
173
+ self.add_zero_attn = add_zero_attn
174
+
175
+ def _reset_parameters(self):
176
+ if self._qkv_same_embed_dim:
177
+ xavier_uniform_(self.in_proj_weight)
178
+ else:
179
+ xavier_uniform_(self.q_proj_weight)
180
+ xavier_uniform_(self.k_proj_weight)
181
+ xavier_uniform_(self.v_proj_weight)
182
+
183
+ if self.in_proj_bias is not None:
184
+ constant_(self.in_proj_bias, 0.0)
185
+ constant_(self.out_proj.bias, 0.0)
186
+
187
+ if self.bias_k is not None:
188
+ xavier_normal_(self.bias_k)
189
+ if self.bias_v is not None:
190
+ xavier_normal_(self.bias_v)
191
+
192
+ def __setstate__(self, state):
193
+ # Support loading old MultiheadAttention checkpoints generated by v1.1.0
194
+ if "_qkv_same_embed_dim" not in state:
195
+ state["_qkv_same_embed_dim"] = True
196
+
197
+ super(MultiheadAttention, self).__setstate__(state)
198
+
199
+ def forward(
200
+ self,
201
+ query: Tensor,
202
+ key: Tensor,
203
+ value: Tensor,
204
+ key_padding_mask: Optional[Tensor] = None,
205
+ need_weights: bool = True,
206
+ attn_mask: Optional[Tensor] = None,
207
+ average_attn_weights: bool = True,
208
+ cache=None,
209
+ ) -> Tuple[Tensor, Optional[Tensor]]:
210
+ r"""
211
+ Args:
212
+ query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
213
+ or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
214
+ :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
215
+ Queries are compared against key-value pairs to produce the output.
216
+ See "Attention Is All You Need" for more details.
217
+ key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
218
+ or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
219
+ :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
220
+ See "Attention Is All You Need" for more details.
221
+ value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
222
+ ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
223
+ sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
224
+ See "Attention Is All You Need" for more details.
225
+ key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
226
+ to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
227
+ Binary and byte masks are supported.
228
+ For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
229
+ the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
230
+ need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
231
+ Default: ``True``.
232
+ attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
233
+ :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
234
+ :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
235
+ broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
236
+ Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
237
+ corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
238
+ corresponding position is not allowed to attend. For a float mask, the mask values will be added to
239
+ the attention weight.
240
+ average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
241
+ heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
242
+ effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
243
+
244
+ Outputs:
245
+ - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
246
+ :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
247
+ where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
248
+ embedding dimension ``embed_dim``.
249
+ - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
250
+ returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
251
+ :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
252
+ :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
253
+ head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
254
+
255
+ .. note::
256
+ `batch_first` argument is ignored for unbatched inputs.
257
+ """
258
+ is_batched = query.dim() == 3
259
+ if key_padding_mask is not None:
260
+ _kpm_dtype = key_padding_mask.dtype
261
+ if _kpm_dtype != torch.bool and not torch.is_floating_point(
262
+ key_padding_mask
263
+ ):
264
+ raise AssertionError(
265
+ "only bool and floating types of key_padding_mask are supported"
266
+ )
267
+ why_not_fast_path = ""
268
+ if not is_batched:
269
+ why_not_fast_path = (
270
+ f"input not batched; expected query.dim() of 3 but got {query.dim()}"
271
+ )
272
+ elif query is not key or key is not value:
273
+ # When lifting this restriction, don't forget to either
274
+ # enforce that the dtypes all match or test cases where
275
+ # they don't!
276
+ why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
277
+ elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:
278
+ why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
279
+ elif (
280
+ self.in_proj_weight is not None and query.dtype != self.in_proj_weight.dtype
281
+ ):
282
+ # this case will fail anyway, but at least they'll get a useful error message.
283
+ why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
284
+ elif self.training:
285
+ why_not_fast_path = "training is enabled"
286
+ elif not self.batch_first:
287
+ why_not_fast_path = "batch_first was not True"
288
+ elif self.bias_k is not None:
289
+ why_not_fast_path = "self.bias_k was not None"
290
+ elif self.bias_v is not None:
291
+ why_not_fast_path = "self.bias_v was not None"
292
+ elif self.dropout:
293
+ why_not_fast_path = f"dropout was {self.dropout}, required zero"
294
+ elif self.add_zero_attn:
295
+ why_not_fast_path = "add_zero_attn was enabled"
296
+ elif not self._qkv_same_embed_dim:
297
+ why_not_fast_path = "_qkv_same_embed_dim was not True"
298
+ elif attn_mask is not None:
299
+ why_not_fast_path = "attn_mask was not None"
300
+ elif query.is_nested and key_padding_mask is not None:
301
+ why_not_fast_path = (
302
+ "key_padding_mask is not supported with NestedTensor input"
303
+ )
304
+ elif self.num_heads % 2 == 1:
305
+ why_not_fast_path = "num_heads is odd"
306
+ elif torch.is_autocast_enabled():
307
+ why_not_fast_path = "autocast is enabled"
308
+
309
+ if not why_not_fast_path:
310
+ tensor_args = (
311
+ query,
312
+ key,
313
+ value,
314
+ self.in_proj_weight,
315
+ self.in_proj_bias,
316
+ self.out_proj.weight,
317
+ self.out_proj.bias,
318
+ )
319
+ # We have to use list comprehensions below because TorchScript does not support
320
+ # generator expressions.
321
+ if torch.overrides.has_torch_function(tensor_args):
322
+ why_not_fast_path = "some Tensor argument has_torch_function"
323
+ elif not all(
324
+ [
325
+ (x is None or x.is_cuda or "cpu" in str(x.device))
326
+ for x in tensor_args
327
+ ]
328
+ ):
329
+ why_not_fast_path = "some Tensor argument is neither CUDA nor CPU"
330
+ elif torch.is_grad_enabled() and any(
331
+ [x is not None and x.requires_grad for x in tensor_args]
332
+ ):
333
+ why_not_fast_path = (
334
+ "grad is enabled and at least one of query or the "
335
+ "input/output projection weights or biases requires_grad"
336
+ )
337
+ if not why_not_fast_path:
338
+ return torch._native_multi_head_attention(
339
+ query,
340
+ key,
341
+ value,
342
+ self.embed_dim,
343
+ self.num_heads,
344
+ self.in_proj_weight,
345
+ self.in_proj_bias,
346
+ self.out_proj.weight,
347
+ self.out_proj.bias,
348
+ key_padding_mask if key_padding_mask is not None else attn_mask,
349
+ need_weights,
350
+ average_attn_weights,
351
+ 1
352
+ if key_padding_mask is not None
353
+ else 0
354
+ if attn_mask is not None
355
+ else None,
356
+ )
357
+
358
+ any_nested = query.is_nested or key.is_nested or value.is_nested
359
+ assert not any_nested, (
360
+ "MultiheadAttention does not support NestedTensor outside of its fast path. "
361
+ + f"The fast path was not hit because {why_not_fast_path}"
362
+ )
363
+
364
+ if self.batch_first and is_batched:
365
+ # make sure that the transpose op does not affect the "is" property
366
+ if key is value:
367
+ if query is key:
368
+ query = key = value = query.transpose(1, 0)
369
+ else:
370
+ query, key = [x.transpose(1, 0) for x in (query, key)]
371
+ value = key
372
+ else:
373
+ query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
374
+
375
+ if not self._qkv_same_embed_dim:
376
+ attn_output, attn_output_weights = multi_head_attention_forward_patched(
377
+ query,
378
+ key,
379
+ value,
380
+ self.embed_dim,
381
+ self.num_heads,
382
+ self.in_proj_weight,
383
+ self.in_proj_bias,
384
+ self.bias_k,
385
+ self.bias_v,
386
+ self.add_zero_attn,
387
+ self.dropout,
388
+ self.out_proj.weight,
389
+ self.out_proj.bias,
390
+ training=self.training,
391
+ key_padding_mask=key_padding_mask,
392
+ need_weights=need_weights,
393
+ attn_mask=attn_mask,
394
+ use_separate_proj_weight=True,
395
+ q_proj_weight=self.q_proj_weight,
396
+ k_proj_weight=self.k_proj_weight,
397
+ v_proj_weight=self.v_proj_weight,
398
+ average_attn_weights=average_attn_weights,
399
+ cache=cache,
400
+ )
401
+ else:
402
+ attn_output, attn_output_weights = multi_head_attention_forward_patched(
403
+ query,
404
+ key,
405
+ value,
406
+ self.embed_dim,
407
+ self.num_heads,
408
+ self.in_proj_weight,
409
+ self.in_proj_bias,
410
+ self.bias_k,
411
+ self.bias_v,
412
+ self.add_zero_attn,
413
+ self.dropout,
414
+ self.out_proj.weight,
415
+ self.out_proj.bias,
416
+ training=self.training,
417
+ key_padding_mask=key_padding_mask,
418
+ need_weights=need_weights,
419
+ attn_mask=attn_mask,
420
+ average_attn_weights=average_attn_weights,
421
+ cache=cache,
422
+ )
423
+ if self.batch_first and is_batched:
424
+ return attn_output.transpose(1, 0), attn_output_weights
425
+ else:
426
+ return attn_output, attn_output_weights
lipsync-models/GPT_SoVITS/AR/modules/embedding.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py
2
+ import math
3
+
4
+ import torch
5
+ from torch import nn
6
+
7
+
8
+ class TokenEmbedding(nn.Module):
9
+ def __init__(
10
+ self,
11
+ embedding_dim: int,
12
+ vocab_size: int,
13
+ dropout: float = 0.0,
14
+ ):
15
+ super().__init__()
16
+
17
+ self.vocab_size = vocab_size
18
+ self.embedding_dim = embedding_dim
19
+
20
+ self.dropout = torch.nn.Dropout(p=dropout)
21
+ self.word_embeddings = nn.Embedding(
22
+ self.vocab_size, self.embedding_dim)
23
+
24
+ @property
25
+ def weight(self) -> torch.Tensor:
26
+ return self.word_embeddings.weight
27
+
28
+ def embedding(self, index: int) -> torch.Tensor:
29
+ return self.word_embeddings.weight[index: index + 1]
30
+
31
+ def forward(self, x: torch.Tensor):
32
+ x = self.word_embeddings(x)
33
+ x = self.dropout(x)
34
+ return x
35
+
36
+
37
+ class SinePositionalEmbedding(nn.Module):
38
+ def __init__(
39
+ self,
40
+ embedding_dim: int,
41
+ dropout: float = 0.0,
42
+ scale: bool = False,
43
+ alpha: bool = False,
44
+ ):
45
+ super().__init__()
46
+ self.embedding_dim = embedding_dim
47
+ self.x_scale = math.sqrt(embedding_dim) if scale else 1.0
48
+ self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
49
+ self.dropout = torch.nn.Dropout(p=dropout)
50
+
51
+ self.reverse = False
52
+ self.register_buffer('pe', self.calc_pe(4000), persistent=False)
53
+
54
+ def calc_pe(self, size: int) -> torch.Tensor:
55
+ pe = torch.zeros(size, self.embedding_dim)
56
+ if self.reverse:
57
+ position = torch.arange(
58
+ size - 1, -1, -1.0, dtype=torch.float32
59
+ ).unsqueeze(1)
60
+ else:
61
+ position = torch.arange(
62
+ 0, size, dtype=torch.float32).unsqueeze(1)
63
+ div_term = torch.exp(
64
+ torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
65
+ * -(math.log(10000.0) / self.embedding_dim)
66
+ )
67
+ pe[:, 0::2] = torch.sin(position * div_term)
68
+ pe[:, 1::2] = torch.cos(position * div_term)
69
+ pe = pe.unsqueeze(0)
70
+
71
+ return pe
72
+
73
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
74
+ output = x.unsqueeze(-1) if x.ndim == 2 else x
75
+ output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
76
+ return self.dropout(output)
lipsync-models/GPT_SoVITS/AR/modules/lr_schedulers.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/lr_schedulers.py
2
+ import math
3
+
4
+ import torch
5
+ from matplotlib import pyplot as plt
6
+ from torch import nn
7
+ from torch.optim import Adam
8
+
9
+
10
+ class WarmupCosineLRSchedule(torch.optim.lr_scheduler._LRScheduler):
11
+ """
12
+ Implements Warmup learning rate schedule until 'warmup_steps', going from 'init_lr' to 'peak_lr' for multiple optimizers.
13
+ """
14
+
15
+ def __init__(
16
+ self,
17
+ optimizer,
18
+ init_lr,
19
+ peak_lr,
20
+ end_lr,
21
+ warmup_steps=10000,
22
+ total_steps=400000,
23
+ current_step=0,
24
+ ):
25
+ self.init_lr = init_lr
26
+ self.peak_lr = peak_lr
27
+ self.end_lr = end_lr
28
+ self.optimizer = optimizer
29
+ self._warmup_rate = (peak_lr - init_lr) / warmup_steps
30
+ self._decay_rate = (end_lr - peak_lr) / (total_steps - warmup_steps)
31
+ self._current_step = current_step
32
+ self.lr = init_lr
33
+ self.warmup_steps = warmup_steps
34
+ self.total_steps = total_steps
35
+ self._last_lr = [self.lr]
36
+
37
+ def set_lr(self, lr):
38
+ self._last_lr = [g["lr"] for g in self.optimizer.param_groups]
39
+ for g in self.optimizer.param_groups:
40
+ # g['lr'] = lr
41
+ g["lr"] = self.end_lr ###锁定用线性
42
+
43
+ def step(self):
44
+ if self._current_step < self.warmup_steps:
45
+ lr = self.init_lr + self._warmup_rate * self._current_step
46
+
47
+ elif self._current_step > self.total_steps:
48
+ lr = self.end_lr
49
+
50
+ else:
51
+ decay_ratio = (self._current_step - self.warmup_steps) / (
52
+ self.total_steps - self.warmup_steps
53
+ )
54
+ if decay_ratio < 0.0 or decay_ratio > 1.0:
55
+ raise RuntimeError(
56
+ "Decay ratio must be in [0.0, 1.0]. Fix LR scheduler settings."
57
+ )
58
+ coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
59
+ lr = self.end_lr + coeff * (self.peak_lr - self.end_lr)
60
+
61
+ self.lr = lr = self.end_lr = 0.002 ###锁定用线性###不听话,直接锁定!
62
+ self.set_lr(lr)
63
+ self.lr = lr
64
+ self._current_step += 1
65
+ return self.lr
66
+
67
+
68
+ if __name__ == "__main__":
69
+ m = nn.Linear(10, 10)
70
+ opt = Adam(m.parameters(), lr=1e-4)
71
+ s = WarmupCosineLRSchedule(
72
+ opt, 1e-6, 2e-4, 1e-6, warmup_steps=2000, total_steps=20000, current_step=0
73
+ )
74
+ lrs = []
75
+ for i in range(25000):
76
+ s.step()
77
+ lrs.append(s.lr)
78
+ print(s.lr)
79
+
80
+ plt.plot(lrs)
81
+ plt.plot(range(0, 25000), lrs)
82
+ plt.show()
lipsync-models/GPT_SoVITS/AR/modules/optim.py ADDED
@@ -0,0 +1,622 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
2
+ #
3
+ # See ../LICENSE for clarification regarding multiple authors
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ import contextlib
17
+ import logging
18
+ from collections import defaultdict
19
+ from typing import List
20
+ from typing import Tuple
21
+
22
+ import torch
23
+ from torch import Tensor
24
+ from torch.optim import Optimizer
25
+
26
+
27
+ class BatchedOptimizer(Optimizer):
28
+ """
29
+ This class adds to class Optimizer the capability to optimize parameters in batches:
30
+ it will stack the parameters and their grads for you so the optimizer can work
31
+ on tensors with an extra leading dimension. This is intended for speed with GPUs,
32
+ as it reduces the number of kernels launched in the optimizer.
33
+
34
+ Args:
35
+ params:
36
+ """
37
+
38
+ def __init__(self, params, defaults):
39
+ super(BatchedOptimizer, self).__init__(params, defaults)
40
+
41
+ @contextlib.contextmanager
42
+ def batched_params(self, param_group, group_params_names):
43
+ """
44
+ This function returns (technically, yields) a list of
45
+ of tuples (p, state), where
46
+ p is a `fake` parameter that is stacked (over axis 0) from real parameters
47
+ that share the same shape, and its gradient is also stacked;
48
+ `state` is the state corresponding to this batch of parameters
49
+ (it will be physically located in the "state" for one of the real
50
+ parameters, the last one that has any particular shape and dtype).
51
+
52
+ This function is decorated as a context manager so that it can
53
+ write parameters back to their "real" locations.
54
+
55
+ The idea is, instead of doing:
56
+ <code>
57
+ for p in group["params"]:
58
+ state = self.state[p]
59
+ ...
60
+ </code>
61
+ you can do:
62
+ <code>
63
+ with self.batched_params(group["params"]) as batches:
64
+ for p, state, p_names in batches:
65
+ ...
66
+ </code>
67
+
68
+ Args:
69
+ group: a parameter group, which is a list of parameters; should be
70
+ one of self.param_groups.
71
+ group_params_names: name for each parameter in group,
72
+ which is List[str].
73
+ """
74
+ batches = defaultdict(
75
+ list
76
+ ) # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
77
+ batches_names = defaultdict(
78
+ list
79
+ ) # `batches` maps from tuple (dtype_as_str,*shape) to list of str
80
+
81
+ assert len(param_group) == len(group_params_names)
82
+ for p, named_p in zip(param_group, group_params_names):
83
+ key = (str(p.dtype), *p.shape)
84
+ batches[key].append(p)
85
+ batches_names[key].append(named_p)
86
+
87
+ batches_names_keys = list(batches_names.keys())
88
+ sorted_idx = sorted(
89
+ range(len(batches_names)), key=lambda i: batches_names_keys[i])
90
+ batches_names = [
91
+ batches_names[batches_names_keys[idx]] for idx in sorted_idx
92
+ ]
93
+ batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
94
+
95
+ stacked_params_dict = dict()
96
+
97
+ # turn batches into a list, in deterministic order.
98
+ # tuples will contain tuples of (stacked_param, state, stacked_params_names),
99
+ # one for each batch in `batches`.
100
+ tuples = []
101
+
102
+ for batch, batch_names in zip(batches, batches_names):
103
+ p = batch[0]
104
+ # we arbitrarily store the state in the
105
+ # state corresponding to the 1st parameter in the
106
+ # group. class Optimizer will take care of saving/loading state.
107
+ state = self.state[p]
108
+ p_stacked = torch.stack(batch)
109
+ grad = torch.stack([
110
+ torch.zeros_like(p) if p.grad is None else p.grad for p in batch
111
+ ])
112
+ p_stacked.grad = grad
113
+ stacked_params_dict[key] = p_stacked
114
+ tuples.append((p_stacked, state, batch_names))
115
+
116
+ yield tuples # <-- calling code will do the actual optimization here!
117
+
118
+ for ((stacked_params, _state, _names), batch) in zip(tuples, batches):
119
+ for i, p in enumerate(batch): # batch is list of Parameter
120
+ p.copy_(stacked_params[i])
121
+
122
+
123
+ class ScaledAdam(BatchedOptimizer):
124
+ """
125
+ Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
126
+ proportional to the norm of that parameter; and also learn the scale of the parameter,
127
+ in log space, subject to upper and lower limits (as if we had factored each parameter as
128
+ param = underlying_param * log_scale.exp())
129
+
130
+
131
+ Args:
132
+ params: The parameters or param_groups to optimize (like other Optimizer subclasses)
133
+ lr: The learning rate. We will typically use a learning rate schedule that starts
134
+ at 0.03 and decreases over time, i.e. much higher than other common
135
+ optimizers.
136
+ clipping_scale: (e.g. 2.0)
137
+ A scale for gradient-clipping: if specified, the normalized gradients
138
+ over the whole model will be clipped to have 2-norm equal to
139
+ `clipping_scale` times the median 2-norm over the most recent period
140
+ of `clipping_update_period` minibatches. By "normalized gradients",
141
+ we mean after multiplying by the rms parameter value for this tensor
142
+ [for non-scalars]; this is appropriate because our update is scaled
143
+ by this quantity.
144
+ betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.
145
+ Must satisfy 0 < beta <= beta2 < 1.
146
+ scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
147
+ scale of each parameter tensor and scalar parameters of the mode..
148
+ If each parameter were decomposed
149
+ as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale
150
+ would be a the scaling factor on the learning rate of p_scale.
151
+ eps: A general-purpose epsilon to prevent division by zero
152
+ param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
153
+ learning the scale on the parameters (we'll constrain the rms of each non-scalar
154
+ parameter tensor to be >= this value)
155
+ param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
156
+ learning the scale on the parameters (we'll constrain the rms of each non-scalar
157
+ parameter tensor to be <= this value)
158
+ scalar_max: Maximum absolute value for scalar parameters (applicable if your
159
+ model has any parameters with numel() == 1).
160
+ size_update_period: The periodicity, in steps, with which we update the size (scale)
161
+ of the parameter tensor. This is provided to save a little time
162
+ in the update.
163
+ clipping_update_period: if clipping_scale is specified, this is the period
164
+ """
165
+
166
+ def __init__(
167
+ self,
168
+ params,
169
+ lr=3e-02,
170
+ clipping_scale=None,
171
+ betas=(0.9, 0.98),
172
+ scalar_lr_scale=0.1,
173
+ eps=1.0e-08,
174
+ param_min_rms=1.0e-05,
175
+ param_max_rms=3.0,
176
+ scalar_max=10.0,
177
+ size_update_period=4,
178
+ clipping_update_period=100,
179
+ parameters_names=None,
180
+ show_dominant_parameters=True, ):
181
+
182
+ assert parameters_names is not None, (
183
+ "Please prepare parameters_names,"
184
+ "which is a List[List[str]]. Each List[str] is for a group"
185
+ "and each str is for a parameter")
186
+ defaults = dict(
187
+ lr=lr,
188
+ clipping_scale=clipping_scale,
189
+ betas=betas,
190
+ scalar_lr_scale=scalar_lr_scale,
191
+ eps=eps,
192
+ param_min_rms=param_min_rms,
193
+ param_max_rms=param_max_rms,
194
+ scalar_max=scalar_max,
195
+ size_update_period=size_update_period,
196
+ clipping_update_period=clipping_update_period, )
197
+
198
+ super(ScaledAdam, self).__init__(params, defaults)
199
+ assert len(self.param_groups) == len(parameters_names)
200
+ self.parameters_names = parameters_names
201
+ self.show_dominant_parameters = show_dominant_parameters
202
+
203
+ def __setstate__(self, state):
204
+ super(ScaledAdam, self).__setstate__(state)
205
+
206
+ @torch.no_grad()
207
+ def step(self, closure=None):
208
+ """Performs a single optimization step.
209
+
210
+ Arguments:
211
+ closure (callable, optional): A closure that reevaluates the model
212
+ and returns the loss.
213
+ """
214
+ loss = None
215
+ if closure is not None:
216
+ with torch.enable_grad():
217
+ loss = closure()
218
+
219
+ batch = True
220
+
221
+ for group, group_params_names in zip(self.param_groups,
222
+ self.parameters_names):
223
+
224
+ with self.batched_params(group["params"],
225
+ group_params_names) as batches:
226
+
227
+ # batches is list of pairs (stacked_param, state). stacked_param is like
228
+ # a regular parameter, and will have a .grad, but the 1st dim corresponds to
229
+ # a stacking dim, it is not a real dim.
230
+
231
+ if (len(batches[0][1]) ==
232
+ 0): # if len(first state) == 0: not yet initialized
233
+ clipping_scale = 1
234
+ else:
235
+ clipping_scale = self._get_clipping_scale(group, batches)
236
+
237
+ for p, state, _ in batches:
238
+ # Perform optimization step.
239
+ # grad is not going to be None, we handled that when creating the batches.
240
+ grad = p.grad
241
+ if grad.is_sparse:
242
+ raise RuntimeError(
243
+ "ScaledAdam optimizer does not support sparse gradients"
244
+ )
245
+ # State initialization
246
+ if len(state) == 0:
247
+ self._init_state(group, p, state)
248
+
249
+ self._step_one_batch(group, p, state, clipping_scale)
250
+
251
+ return loss
252
+
253
+ def _init_state(self, group: dict, p: Tensor, state: dict):
254
+ """
255
+ Initializes state dict for parameter 'p'. Assumes that dim 0 of tensor p
256
+ is actually the batch dimension, corresponding to batched-together
257
+ parameters of a given shape.
258
+
259
+
260
+ Args:
261
+ group: Dict to look up configuration values.
262
+ p: The parameter that we are initializing the state for
263
+ state: Dict from string to whatever state we are initializing
264
+ """
265
+ size_update_period = group["size_update_period"]
266
+
267
+ state["step"] = 0
268
+
269
+ kwargs = {"device": p.device, "dtype": p.dtype}
270
+
271
+ # 'delta' implements conventional momentum. There are
272
+ # several different kinds of update going on, so rather than
273
+ # compute "exp_avg" like in Adam, we store and decay a
274
+ # parameter-change "delta", which combines all forms of
275
+ # update. this is equivalent to how it's done in Adam,
276
+ # except for the first few steps.
277
+ state["delta"] = torch.zeros_like(
278
+ p, memory_format=torch.preserve_format)
279
+
280
+ batch_size = p.shape[0]
281
+ numel = p.numel() // batch_size
282
+ numel = p.numel()
283
+
284
+ if numel > 1:
285
+ # "param_rms" just periodically records the scalar root-mean-square value of
286
+ # the parameter tensor.
287
+ # it has a shape like (batch_size, 1, 1, 1, 1)
288
+ param_rms = (
289
+ (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt())
290
+ state["param_rms"] = param_rms
291
+
292
+ state["scale_exp_avg_sq"] = torch.zeros_like(param_rms)
293
+ state["scale_grads"] = torch.zeros(size_update_period,
294
+ *param_rms.shape, **kwargs)
295
+
296
+ # exp_avg_sq is the weighted sum of scaled gradients. as in Adam.
297
+ state["exp_avg_sq"] = torch.zeros_like(
298
+ p, memory_format=torch.preserve_format)
299
+
300
+ def _get_clipping_scale(self,
301
+ group: dict,
302
+ tuples: List[Tuple[Tensor, dict, List[str]]]
303
+ ) -> float:
304
+ """
305
+ Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients
306
+ by this amount before applying the rest of the update.
307
+
308
+ Args:
309
+ group: the parameter group, an item in self.param_groups
310
+ tuples: a list of tuples of (param, state, param_names)
311
+ where param is a batched set of parameters,
312
+ with a .grad (1st dim is batch dim)
313
+ and state is the state-dict where optimization parameters are kept.
314
+ param_names is a List[str] while each str is name for a parameter
315
+ in batched set of parameters "param".
316
+ """
317
+ assert len(tuples) >= 1
318
+ clipping_scale = group["clipping_scale"]
319
+ (first_p, first_state, _) = tuples[0]
320
+ step = first_state["step"]
321
+ if clipping_scale is None or step == 0:
322
+ # no clipping. return early on step == 0 because the other
323
+ # parameters' state won't have been initialized yet.
324
+ return 1.0
325
+ clipping_update_period = group["clipping_update_period"]
326
+
327
+ tot_sumsq = torch.tensor(0.0, device=first_p.device)
328
+ for (p, state, param_names) in tuples:
329
+ grad = p.grad
330
+ if grad.is_sparse:
331
+ raise RuntimeError(
332
+ "ScaledAdam optimizer does not support sparse gradients")
333
+ if p.numel() == p.shape[0]: # a batch of scalars
334
+ tot_sumsq += (grad**2).sum() # sum() to change shape [1] to []
335
+ else:
336
+ tot_sumsq += ((grad * state["param_rms"])**2).sum()
337
+
338
+ tot_norm = tot_sumsq.sqrt()
339
+ if "model_norms" not in first_state:
340
+ first_state["model_norms"] = torch.zeros(
341
+ clipping_update_period, device=p.device)
342
+ first_state["model_norms"][step % clipping_update_period] = tot_norm
343
+
344
+ if step % clipping_update_period == 0:
345
+ # Print some stats.
346
+ # We don't reach here if step == 0 because we would have returned
347
+ # above.
348
+ sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
349
+ quartiles = []
350
+ for n in range(0, 5):
351
+ index = min(
352
+ clipping_update_period - 1,
353
+ (clipping_update_period // 4) * n, )
354
+ quartiles.append(sorted_norms[index].item())
355
+
356
+ median = quartiles[2]
357
+ threshold = clipping_scale * median
358
+ first_state["model_norm_threshold"] = threshold
359
+ percent_clipped = (first_state["num_clipped"] * 100.0 /
360
+ clipping_update_period
361
+ if "num_clipped" in first_state else 0.0)
362
+ first_state["num_clipped"] = 0
363
+ quartiles = " ".join(["%.3e" % x for x in quartiles])
364
+ logging.info(
365
+ f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
366
+ f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
367
+ )
368
+
369
+ if step < clipping_update_period:
370
+ return 1.0 # We have not yet estimated a norm to clip to.
371
+ else:
372
+ try:
373
+ model_norm_threshold = first_state["model_norm_threshold"]
374
+ except KeyError:
375
+ logging.info(
376
+ "Warning: model_norm_threshold not in state: possibly "
377
+ "you changed config when restarting, adding clipping_scale option?"
378
+ )
379
+ return 1.0
380
+ ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
381
+ if ans < 1.0:
382
+ first_state["num_clipped"] += 1
383
+ if ans < 0.1:
384
+ logging.warn(
385
+ f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}"
386
+ )
387
+ if self.show_dominant_parameters:
388
+ assert p.shape[0] == len(param_names)
389
+ self._show_gradient_dominating_parameter(tuples, tot_sumsq)
390
+ return ans
391
+
392
+ def _show_gradient_dominating_parameter(
393
+ self, tuples: List[Tuple[Tensor, dict, List[str]]],
394
+ tot_sumsq: Tensor):
395
+ """
396
+ Show information of parameter wihch dominanting tot_sumsq.
397
+
398
+ Args:
399
+ tuples: a list of tuples of (param, state, param_names)
400
+ where param is a batched set of parameters,
401
+ with a .grad (1st dim is batch dim)
402
+ and state is the state-dict where optimization parameters are kept.
403
+ param_names is a List[str] while each str is name for a parameter
404
+ in batched set of parameters "param".
405
+ tot_sumsq: sumsq of all parameters. Though it's could be calculated
406
+ from tuples, we still pass it to save some time.
407
+ """
408
+ all_sumsq_orig = {}
409
+ for (p, state, batch_param_names) in tuples:
410
+ # p is a stacked batch parameters.
411
+ batch_grad = p.grad
412
+ if p.numel() == p.shape[0]: # a batch of scalars
413
+ batch_sumsq_orig = batch_grad**2
414
+ # Dummpy values used by following `zip` statement.
415
+ batch_rms_orig = torch.ones(p.shape[0])
416
+ else:
417
+ batch_rms_orig = state["param_rms"]
418
+ batch_sumsq_orig = ((batch_grad * batch_rms_orig)**2).sum(
419
+ dim=list(range(1, batch_grad.ndim)))
420
+
421
+ for name, sumsq_orig, rms, grad in zip(batch_param_names,
422
+ batch_sumsq_orig,
423
+ batch_rms_orig, batch_grad):
424
+
425
+ proportion_orig = sumsq_orig / tot_sumsq
426
+ all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
427
+
428
+ assert torch.isclose(
429
+ sum([value[0] for value in all_sumsq_orig.values()]).cpu(),
430
+ torch.tensor(1.0), )
431
+ sorted_by_proportion = {
432
+ k: v
433
+ for k, v in sorted(
434
+ all_sumsq_orig.items(),
435
+ key=lambda item: item[1][0],
436
+ reverse=True, )
437
+ }
438
+ dominant_param_name = next(iter(sorted_by_proportion))
439
+ (dominant_proportion, dominant_sumsq, dominant_rms,
440
+ dominant_grad, ) = sorted_by_proportion[dominant_param_name]
441
+ logging.info(f"Parameter Dominanting tot_sumsq {dominant_param_name}"
442
+ f" with proportion {dominant_proportion:.2f},"
443
+ f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
444
+ f"={dominant_sumsq:.3e},"
445
+ f" grad_sumsq = {(dominant_grad**2).sum():.3e},"
446
+ f" orig_rms_sq={(dominant_rms**2).item():.3e}")
447
+
448
+ def _step_one_batch(self,
449
+ group: dict,
450
+ p: Tensor,
451
+ state: dict,
452
+ clipping_scale: float):
453
+ """
454
+ Do the step for one parameter, which is actually going to be a batch of
455
+ `real` parameters, with dim 0 as the batch dim.
456
+ Args:
457
+ group: dict to look up configuration values
458
+ p: parameter to update (actually multiple parameters stacked together
459
+ as a batch)
460
+ state: state-dict for p, to look up the optimizer state
461
+ """
462
+ lr = group["lr"]
463
+ size_update_period = group["size_update_period"]
464
+ beta1 = group["betas"][0]
465
+
466
+ grad = p.grad
467
+ if clipping_scale != 1.0:
468
+ grad = grad * clipping_scale
469
+ step = state["step"]
470
+ delta = state["delta"]
471
+
472
+ delta.mul_(beta1)
473
+ batch_size = p.shape[0]
474
+ numel = p.numel() // batch_size
475
+ if numel > 1:
476
+ # Update the size/scale of p, and set param_rms
477
+ scale_grads = state["scale_grads"]
478
+ scale_grads[step % size_update_period] = (p * grad).sum(
479
+ dim=list(range(1, p.ndim)), keepdim=True)
480
+ if step % size_update_period == size_update_period - 1:
481
+ param_rms = state["param_rms"] # shape: (batch_size, 1, 1, ..)
482
+ param_rms.copy_((p**2)
483
+ .mean(dim=list(range(1, p.ndim)), keepdim=True)
484
+ .sqrt())
485
+ if step > 0:
486
+ # self._size_update() learns the overall scale on the
487
+ # parameter, by shrinking or expanding it.
488
+ self._size_update(group, scale_grads, p, state)
489
+
490
+ if numel == 1:
491
+ # For parameters with 1 element we just use regular Adam.
492
+ # Updates delta.
493
+ self._step_scalar(group, p, state)
494
+ else:
495
+ self._step(group, p, state)
496
+
497
+ state["step"] = step + 1
498
+
499
+ def _size_update(self,
500
+ group: dict,
501
+ scale_grads: Tensor,
502
+ p: Tensor,
503
+ state: dict) -> None:
504
+ """
505
+ Called only where p.numel() > 1, this updates the scale of the parameter.
506
+ If we imagine: p = underlying_param * scale.exp(), and we are doing
507
+ gradient descent on underlying param and on scale, this function does the update
508
+ on `scale`.
509
+
510
+ Args:
511
+ group: dict to look up configuration values
512
+ scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing
513
+ grads w.r.t. the scales.
514
+ p: The parameter to update
515
+ state: The state-dict of p
516
+ """
517
+
518
+ param_rms = state["param_rms"]
519
+ beta1, beta2 = group["betas"]
520
+ size_lr = group["lr"] * group["scalar_lr_scale"]
521
+ param_min_rms = group["param_min_rms"]
522
+ param_max_rms = group["param_max_rms"]
523
+ eps = group["eps"]
524
+ step = state["step"]
525
+ batch_size = p.shape[0]
526
+
527
+ size_update_period = scale_grads.shape[0]
528
+ # correct beta2 for the size update period: we will have
529
+ # faster decay at this level.
530
+ beta2_corr = beta2**size_update_period
531
+
532
+ scale_exp_avg_sq = state[
533
+ "scale_exp_avg_sq"] # shape: (batch_size, 1, 1, ..)
534
+ scale_exp_avg_sq.mul_(beta2_corr).add_(
535
+ (scale_grads**2).mean(dim=0), # mean over dim `size_update_period`
536
+ alpha=1 - beta2_corr, ) # shape is (batch_size, 1, 1, ...)
537
+
538
+ # The 1st time we reach here is when size_step == 1.
539
+ size_step = (step + 1) // size_update_period
540
+ bias_correction2 = 1 - beta2_corr**size_step
541
+ # we don't bother with bias_correction1; this will help prevent divergence
542
+ # at the start of training.
543
+
544
+ denom = scale_exp_avg_sq.sqrt() + eps
545
+
546
+ scale_step = (-size_lr * (bias_correction2**0.5) *
547
+ scale_grads.sum(dim=0) / denom)
548
+
549
+ is_too_small = param_rms < param_min_rms
550
+ is_too_large = param_rms > param_max_rms
551
+
552
+ # when the param gets too small, just don't shrink it any further.
553
+ scale_step.masked_fill_(is_too_small, 0.0)
554
+ # when it gets too large, stop it from getting any larger.
555
+ scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)
556
+ delta = state["delta"]
557
+ # the factor of (1-beta1) relates to momentum.
558
+ delta.add_(p * scale_step, alpha=(1 - beta1))
559
+
560
+ def _step(self, group: dict, p: Tensor, state: dict):
561
+ """
562
+ This function does the core update of self.step(), in the case where the members of
563
+ the batch have more than 1 element.
564
+
565
+ Args:
566
+ group: A dict which will be used to look up configuration values
567
+ p: The parameter to be updated
568
+ grad: The grad of p
569
+ state: The state-dict corresponding to parameter p
570
+
571
+ This function modifies p.
572
+ """
573
+ grad = p.grad
574
+ lr = group["lr"]
575
+ beta1, beta2 = group["betas"]
576
+ eps = group["eps"]
577
+ param_min_rms = group["param_min_rms"]
578
+ step = state["step"]
579
+
580
+ exp_avg_sq = state["exp_avg_sq"]
581
+ exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1 - beta2))
582
+
583
+ this_step = state["step"] - (state["zero_step"]
584
+ if "zero_step" in state else 0)
585
+ bias_correction2 = 1 - beta2**(this_step + 1)
586
+ if bias_correction2 < 0.99:
587
+ # note: not in-place.
588
+ exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
589
+
590
+ denom = exp_avg_sq.sqrt()
591
+ denom += eps
592
+ grad = grad / denom
593
+
594
+ alpha = -lr * (1 - beta1) * state["param_rms"].clamp(min=param_min_rms)
595
+
596
+ delta = state["delta"]
597
+ delta.add_(grad * alpha)
598
+ p.add_(delta)
599
+
600
+ def _step_scalar(self, group: dict, p: Tensor, state: dict):
601
+ """
602
+ A simplified form of the core update for scalar tensors, where we cannot get a good
603
+ estimate of the parameter rms.
604
+ """
605
+ beta1, beta2 = group["betas"]
606
+ scalar_max = group["scalar_max"]
607
+ eps = group["eps"]
608
+ lr = group["lr"] * group["scalar_lr_scale"]
609
+ grad = p.grad
610
+
611
+ exp_avg_sq = state["exp_avg_sq"] # shape: (batch_size,)
612
+ exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
613
+
614
+ # bias_correction2 is like in Adam. Don't bother with bias_correction1;
615
+ # slower update at the start will help stability anyway.
616
+ bias_correction2 = 1 - beta2**(state["step"] + 1)
617
+ denom = (exp_avg_sq / bias_correction2).sqrt() + eps
618
+
619
+ delta = state["delta"]
620
+ delta.add_(grad / denom, alpha=-lr * (1 - beta1))
621
+ p.clamp_(min=-scalar_max, max=scalar_max)
622
+ p.add_(delta)
lipsync-models/GPT_SoVITS/AR/modules/patched_mha_with_cache.py ADDED
@@ -0,0 +1,472 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch.nn.functional import *
2
+ from torch.nn.functional import (
3
+ _mha_shape_check,
4
+ _canonical_mask,
5
+ _none_or_dtype,
6
+ _in_projection_packed,
7
+ _in_projection,
8
+ )
9
+ from torch.nn import functional as F
10
+ import torch
11
+ # Tensor = torch.Tensor
12
+ # from typing import Callable, List, Optional, Tuple, Union
13
+
14
+
15
+ def multi_head_attention_forward_patched(
16
+ query: Tensor,
17
+ key: Tensor,
18
+ value: Tensor,
19
+ embed_dim_to_check: int,
20
+ num_heads: int,
21
+ in_proj_weight: Optional[Tensor],
22
+ in_proj_bias: Optional[Tensor],
23
+ bias_k: Optional[Tensor],
24
+ bias_v: Optional[Tensor],
25
+ add_zero_attn: bool,
26
+ dropout_p: float,
27
+ out_proj_weight: Tensor,
28
+ out_proj_bias: Optional[Tensor],
29
+ training: bool = True,
30
+ key_padding_mask: Optional[Tensor] = None,
31
+ need_weights: bool = True,
32
+ attn_mask: Optional[Tensor] = None,
33
+ use_separate_proj_weight: bool = False,
34
+ q_proj_weight: Optional[Tensor] = None,
35
+ k_proj_weight: Optional[Tensor] = None,
36
+ v_proj_weight: Optional[Tensor] = None,
37
+ static_k: Optional[Tensor] = None,
38
+ static_v: Optional[Tensor] = None,
39
+ average_attn_weights: bool = True,
40
+ is_causal: bool = False,
41
+ cache=None,
42
+ ) -> Tuple[Tensor, Optional[Tensor]]:
43
+ r"""
44
+ Args:
45
+ query, key, value: map a query and a set of key-value pairs to an output.
46
+ See "Attention Is All You Need" for more details.
47
+ embed_dim_to_check: total dimension of the model.
48
+ num_heads: parallel attention heads.
49
+ in_proj_weight, in_proj_bias: input projection weight and bias.
50
+ bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
51
+ add_zero_attn: add a new batch of zeros to the key and
52
+ value sequences at dim=1.
53
+ dropout_p: probability of an element to be zeroed.
54
+ out_proj_weight, out_proj_bias: the output projection weight and bias.
55
+ training: apply dropout if is ``True``.
56
+ key_padding_mask: if provided, specified padding elements in the key will
57
+ be ignored by the attention. This is an binary mask. When the value is True,
58
+ the corresponding value on the attention layer will be filled with -inf.
59
+ need_weights: output attn_output_weights.
60
+ Default: `True`
61
+ Note: `needs_weight` defaults to `True`, but should be set to `False`
62
+ For best performance when attention weights are not nedeeded.
63
+ *Setting needs_weights to `True`
64
+ leads to a significant performance degradation.*
65
+ attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
66
+ the batches while a 3D mask allows to specify a different mask for the entries of each batch.
67
+ is_causal: If specified, applies a causal mask as attention mask, and ignores
68
+ attn_mask for computing scaled dot product attention.
69
+ Default: ``False``.
70
+ .. warning::
71
+ is_causal is provides a hint that the attn_mask is the
72
+ causal mask.Providing incorrect hints can result in
73
+ incorrect execution, including forward and backward
74
+ compatibility.
75
+ use_separate_proj_weight: the function accept the proj. weights for query, key,
76
+ and value in different forms. If false, in_proj_weight will be used, which is
77
+ a combination of q_proj_weight, k_proj_weight, v_proj_weight.
78
+ q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
79
+ static_k, static_v: static key and value used for attention operators.
80
+ average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.
81
+ Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect
82
+ when ``need_weights=True.``. Default: True
83
+
84
+
85
+ Shape:
86
+ Inputs:
87
+ - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
88
+ the embedding dimension.
89
+ - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
90
+ the embedding dimension.
91
+ - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
92
+ the embedding dimension.
93
+ - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.
94
+ If a FloatTensor is provided, it will be directly added to the value.
95
+ If a BoolTensor is provided, the positions with the
96
+ value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
97
+ - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
98
+ 3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
99
+ S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
100
+ positions. If a BoolTensor is provided, positions with ``True``
101
+ are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
102
+ is provided, it will be added to the attention weight.
103
+ - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
104
+ N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
105
+ - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
106
+ N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
107
+
108
+ Outputs:
109
+ - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
110
+ E is the embedding dimension.
111
+ - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns
112
+ attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
113
+ :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
114
+ :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
115
+ head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.
116
+ """
117
+ tens_ops = (
118
+ query,
119
+ key,
120
+ value,
121
+ in_proj_weight,
122
+ in_proj_bias,
123
+ bias_k,
124
+ bias_v,
125
+ out_proj_weight,
126
+ out_proj_bias,
127
+ )
128
+ if has_torch_function(tens_ops):
129
+ return handle_torch_function(
130
+ multi_head_attention_forward,
131
+ tens_ops,
132
+ query,
133
+ key,
134
+ value,
135
+ embed_dim_to_check,
136
+ num_heads,
137
+ in_proj_weight,
138
+ in_proj_bias,
139
+ bias_k,
140
+ bias_v,
141
+ add_zero_attn,
142
+ dropout_p,
143
+ out_proj_weight,
144
+ out_proj_bias,
145
+ training=training,
146
+ key_padding_mask=key_padding_mask,
147
+ need_weights=need_weights,
148
+ attn_mask=attn_mask,
149
+ is_causal=is_causal,
150
+ use_separate_proj_weight=use_separate_proj_weight,
151
+ q_proj_weight=q_proj_weight,
152
+ k_proj_weight=k_proj_weight,
153
+ v_proj_weight=v_proj_weight,
154
+ static_k=static_k,
155
+ static_v=static_v,
156
+ average_attn_weights=average_attn_weights,
157
+ cache=cache,
158
+ )
159
+
160
+ is_batched = _mha_shape_check(
161
+ query, key, value, key_padding_mask, attn_mask, num_heads
162
+ )
163
+
164
+ # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input
165
+ # is batched, run the computation and before returning squeeze the
166
+ # batch dimension so that the output doesn't carry this temporary batch dimension.
167
+ if not is_batched:
168
+ # unsqueeze if the input is unbatched
169
+ query = query.unsqueeze(1)
170
+ key = key.unsqueeze(1)
171
+ value = value.unsqueeze(1)
172
+ if key_padding_mask is not None:
173
+ key_padding_mask = key_padding_mask.unsqueeze(0)
174
+
175
+ # set up shape vars
176
+ tgt_len, bsz, embed_dim = query.shape
177
+ src_len, _, _ = key.shape
178
+
179
+ key_padding_mask = _canonical_mask(
180
+ mask=key_padding_mask,
181
+ mask_name="key_padding_mask",
182
+ other_type=_none_or_dtype(attn_mask),
183
+ other_name="attn_mask",
184
+ target_type=query.dtype,
185
+ )
186
+
187
+ if is_causal and attn_mask is None:
188
+ raise RuntimeError(
189
+ "Need attn_mask if specifying the is_causal hint. "
190
+ "You may use the Transformer module method "
191
+ "`generate_square_subsequent_mask` to create this mask."
192
+ )
193
+
194
+ if is_causal and key_padding_mask is None and not need_weights:
195
+ # when we have a kpm or need weights, we need attn_mask
196
+ # Otherwise, we use the is_causal hint go as is_causal
197
+ # indicator to SDPA.
198
+ attn_mask = None
199
+ else:
200
+ attn_mask = _canonical_mask(
201
+ mask=attn_mask,
202
+ mask_name="attn_mask",
203
+ other_type=None,
204
+ other_name="",
205
+ target_type=query.dtype,
206
+ check_other=False,
207
+ )
208
+
209
+ if key_padding_mask is not None:
210
+ # We have the attn_mask, and use that to merge kpm into it.
211
+ # Turn off use of is_causal hint, as the merged mask is no
212
+ # longer causal.
213
+ is_causal = False
214
+
215
+ assert (
216
+ embed_dim == embed_dim_to_check
217
+ ), f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
218
+ if isinstance(embed_dim, torch.Tensor):
219
+ # embed_dim can be a tensor when JIT tracing
220
+ head_dim = embed_dim.div(num_heads, rounding_mode="trunc")
221
+ else:
222
+ head_dim = embed_dim // num_heads
223
+ assert (
224
+ head_dim * num_heads == embed_dim
225
+ ), f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
226
+ if use_separate_proj_weight:
227
+ # allow MHA to have different embedding dimensions when separate projection weights are used
228
+ assert (
229
+ key.shape[:2] == value.shape[:2]
230
+ ), f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
231
+ else:
232
+ assert (
233
+ key.shape == value.shape
234
+ ), f"key shape {key.shape} does not match value shape {value.shape}"
235
+
236
+ #
237
+ # compute in-projection
238
+ #
239
+ if not use_separate_proj_weight:
240
+ assert (
241
+ in_proj_weight is not None
242
+ ), "use_separate_proj_weight is False but in_proj_weight is None"
243
+ q, k, v = _in_projection_packed(
244
+ query, key, value, in_proj_weight, in_proj_bias)
245
+ else:
246
+ assert (
247
+ q_proj_weight is not None
248
+ ), "use_separate_proj_weight is True but q_proj_weight is None"
249
+ assert (
250
+ k_proj_weight is not None
251
+ ), "use_separate_proj_weight is True but k_proj_weight is None"
252
+ assert (
253
+ v_proj_weight is not None
254
+ ), "use_separate_proj_weight is True but v_proj_weight is None"
255
+ if in_proj_bias is None:
256
+ b_q = b_k = b_v = None
257
+ else:
258
+ b_q, b_k, b_v = in_proj_bias.chunk(3)
259
+ q, k, v = _in_projection(
260
+ query,
261
+ key,
262
+ value,
263
+ q_proj_weight,
264
+ k_proj_weight,
265
+ v_proj_weight,
266
+ b_q,
267
+ b_k,
268
+ b_v,
269
+ )
270
+
271
+ # The dict cache is not compatible with JIT script mode, so we disable it
272
+ # if cache is not None:
273
+ # if cache["first_infer"] == 1:
274
+ # cache["k"][cache["stage"]] = k
275
+ # # print(0,cache["k"].shape)
276
+ # cache["v"][cache["stage"]] = v
277
+ # else: ###12个layer每个都要留自己的cache_kv
278
+ # # print(1,cache["k"].shape)
279
+ # cache["k"][cache["stage"]] = torch.cat(
280
+ # [cache["k"][cache["stage"]], k], 0
281
+ # ) ##本来时序是1,但是proj的时候可能transpose了所以时序到0维了
282
+ # cache["v"][cache["stage"]] = torch.cat([cache["v"][cache["stage"]], v], 0)
283
+ # # print(2, cache["k"].shape)
284
+ # src_len = cache["k"][cache["stage"]].shape[0]
285
+ # k = cache["k"][cache["stage"]]
286
+ # v = cache["v"][cache["stage"]]
287
+ # # if attn_mask is not None:
288
+ # # attn_mask=attn_mask[-1:,]
289
+ # # print(attn_mask.shape,attn_mask)
290
+ # cache["stage"] = (cache["stage"] + 1) % cache["all_stage"]
291
+ # # print(2333,cache)
292
+ # prep attention mask
293
+
294
+ attn_mask = _canonical_mask(
295
+ mask=attn_mask,
296
+ mask_name="attn_mask",
297
+ other_type=None,
298
+ other_name="",
299
+ target_type=q.dtype,
300
+ check_other=False,
301
+ )
302
+
303
+ if attn_mask is not None:
304
+ # ensure attn_mask's dim is 3
305
+ if attn_mask.dim() == 2:
306
+ correct_2d_size = (tgt_len, src_len)
307
+ if attn_mask.shape != correct_2d_size:
308
+ raise RuntimeError(
309
+ f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}."
310
+ )
311
+ attn_mask = attn_mask.unsqueeze(0)
312
+ elif attn_mask.dim() == 3:
313
+ correct_3d_size = (bsz * num_heads, tgt_len, src_len)
314
+ if attn_mask.shape != correct_3d_size:
315
+ raise RuntimeError(
316
+ f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}."
317
+ )
318
+ else:
319
+ raise RuntimeError(
320
+ f"attn_mask's dimension {attn_mask.dim()} is not supported"
321
+ )
322
+
323
+ # add bias along batch dimension (currently second)
324
+ if bias_k is not None and bias_v is not None:
325
+ assert static_k is None, "bias cannot be added to static key."
326
+ assert static_v is None, "bias cannot be added to static value."
327
+ k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
328
+ v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
329
+ if attn_mask is not None:
330
+ attn_mask = pad(attn_mask, (0, 1))
331
+ if key_padding_mask is not None:
332
+ key_padding_mask = pad(key_padding_mask, (0, 1))
333
+ else:
334
+ assert bias_k is None
335
+ assert bias_v is None
336
+
337
+ #
338
+ # reshape q, k, v for multihead attention and make em batch first
339
+ #
340
+ q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
341
+ if static_k is None:
342
+ k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
343
+ else:
344
+ # TODO finish disentangling control flow so we don't do in-projections when statics are passed
345
+ assert (
346
+ static_k.size(0) == bsz * num_heads
347
+ ), f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
348
+ assert (
349
+ static_k.size(2) == head_dim
350
+ ), f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
351
+ k = static_k
352
+ if static_v is None:
353
+ v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
354
+ else:
355
+ # TODO finish disentangling control flow so we don't do in-projections when statics are passed
356
+ assert (
357
+ static_v.size(0) == bsz * num_heads
358
+ ), f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
359
+ assert (
360
+ static_v.size(2) == head_dim
361
+ ), f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
362
+ v = static_v
363
+
364
+ # add zero attention along batch dimension (now first)
365
+ if add_zero_attn:
366
+ zero_attn_shape = (bsz * num_heads, 1, head_dim)
367
+ k = torch.cat(
368
+ [k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1
369
+ )
370
+ v = torch.cat(
371
+ [v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1
372
+ )
373
+ if attn_mask is not None:
374
+ attn_mask = pad(attn_mask, (0, 1))
375
+ if key_padding_mask is not None:
376
+ key_padding_mask = pad(key_padding_mask, (0, 1))
377
+
378
+ # update source sequence length after adjustments
379
+ src_len = k.size(1)
380
+
381
+ # merge key padding and attention masks
382
+ if key_padding_mask is not None:
383
+ assert key_padding_mask.shape == (
384
+ bsz,
385
+ src_len,
386
+ ), f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
387
+ key_padding_mask = (
388
+ key_padding_mask.view(bsz, 1, 1, src_len)
389
+ .expand(-1, num_heads, -1, -1)
390
+ .reshape(bsz * num_heads, 1, src_len)
391
+ )
392
+ if attn_mask is None:
393
+ attn_mask = key_padding_mask
394
+ else:
395
+ attn_mask = attn_mask + key_padding_mask
396
+
397
+ # adjust dropout probability
398
+ if not training:
399
+ dropout_p = 0.0
400
+
401
+ #
402
+ # (deep breath) calculate attention and out projection
403
+ #
404
+
405
+ if need_weights:
406
+ B, Nt, E = q.shape
407
+ q_scaled = q / math.sqrt(E)
408
+
409
+ assert not (
410
+ is_causal and attn_mask is None
411
+ ), "FIXME: is_causal not implemented for need_weights"
412
+
413
+ if attn_mask is not None:
414
+ attn_output_weights = torch.baddbmm(
415
+ attn_mask, q_scaled, k.transpose(-2, -1)
416
+ )
417
+ else:
418
+ attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
419
+ attn_output_weights = softmax(attn_output_weights, dim=-1)
420
+ if dropout_p > 0.0:
421
+ attn_output_weights = dropout(attn_output_weights, p=dropout_p)
422
+
423
+ attn_output = torch.bmm(attn_output_weights, v)
424
+
425
+ attn_output = (
426
+ attn_output.transpose(0, 1).contiguous().view(
427
+ tgt_len * bsz, embed_dim)
428
+ )
429
+ attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
430
+ attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
431
+
432
+ # optionally average attention weights over heads
433
+ attn_output_weights = attn_output_weights.view(
434
+ bsz, num_heads, tgt_len, src_len)
435
+ if average_attn_weights:
436
+ attn_output_weights = attn_output_weights.mean(dim=1)
437
+
438
+ if not is_batched:
439
+ # squeeze the output if input was unbatched
440
+ attn_output = attn_output.squeeze(1)
441
+ attn_output_weights = attn_output_weights.squeeze(0)
442
+ return attn_output, attn_output_weights
443
+ else:
444
+ # attn_mask can be either (L,S) or (N*num_heads, L, S)
445
+ # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
446
+ # in order to match the input for SDPA of (N, num_heads, L, S)
447
+ if attn_mask is not None:
448
+ if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
449
+ attn_mask = attn_mask.unsqueeze(0)
450
+ else:
451
+ attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
452
+
453
+ q = q.view(bsz, num_heads, tgt_len, head_dim)
454
+ k = k.view(bsz, num_heads, src_len, head_dim)
455
+ v = v.view(bsz, num_heads, src_len, head_dim)
456
+
457
+ # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
458
+ attn_output = scaled_dot_product_attention(
459
+ q, k, v, attn_mask, dropout_p, is_causal
460
+ )
461
+
462
+ attn_output = (
463
+ attn_output.permute(2, 0, 1, 3).contiguous().view(
464
+ bsz * tgt_len, embed_dim)
465
+ )
466
+
467
+ attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
468
+ attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
469
+ if not is_batched:
470
+ # squeeze the output if input was unbatched
471
+ attn_output = attn_output.squeeze(1)
472
+ return attn_output, None
lipsync-models/GPT_SoVITS/AR/modules/scaling.py ADDED
@@ -0,0 +1,335 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
2
+ #
3
+ # See ../../../../LICENSE for clarification regarding multiple authors
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ import logging
17
+ import math
18
+ import random
19
+ from typing import Optional
20
+ from typing import Tuple
21
+ from typing import Union
22
+
23
+ import torch
24
+ import torch.nn as nn
25
+ from torch import Tensor
26
+
27
+
28
+ class DoubleSwishFunction(torch.autograd.Function):
29
+ """
30
+ double_swish(x) = x * torch.sigmoid(x-1)
31
+ This is a definition, originally motivated by its close numerical
32
+ similarity to swish(swish(x)), where swish(x) = x * sigmoid(x).
33
+
34
+ Memory-efficient derivative computation:
35
+ double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
36
+ double_swish'(x) = d/dx double_swish(x) = x * s'(x) + x' * s(x) = x * s'(x) + s(x).
37
+ Now, s'(x) = s(x) * (1-s(x)).
38
+ double_swish'(x) = x * s'(x) + s(x).
39
+ = x * s(x) * (1-s(x)) + s(x).
40
+ = double_swish(x) * (1-s(x)) + s(x)
41
+ ... so we just need to remember s(x) but not x itself.
42
+ """
43
+
44
+ @staticmethod
45
+ def forward(ctx, x: Tensor) -> Tensor:
46
+ requires_grad = x.requires_grad
47
+ x_dtype = x.dtype
48
+ if x.dtype == torch.float16:
49
+ x = x.to(torch.float32)
50
+
51
+ s = torch.sigmoid(x - 1.0)
52
+ y = x * s
53
+
54
+ if requires_grad:
55
+ deriv = y * (1 - s) + s
56
+ # notes on derivative of x * sigmoid(x - 1):
57
+ # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
58
+ # min \simeq -0.043638. Take floor as -0.043637 so it's a lower bund
59
+ # max \simeq 1.1990. Take ceil to be 1.2 so it's an upper bound.
60
+ # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
61
+ # floors), should be expectation-preserving.
62
+ floor = -0.043637
63
+ ceil = 1.2
64
+ d_scaled = (deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
65
+ deriv
66
+ )
67
+ if __name__ == "__main__":
68
+ # for self-testing only.
69
+ assert d_scaled.min() >= 0.0
70
+ assert d_scaled.max() < 256.0
71
+ d_int = d_scaled.to(torch.uint8)
72
+ ctx.save_for_backward(d_int)
73
+ if x.dtype == torch.float16 or torch.is_autocast_enabled():
74
+ y = y.to(torch.float16)
75
+ return y
76
+
77
+ @staticmethod
78
+ def backward(ctx, y_grad: Tensor) -> Tensor:
79
+ (d,) = ctx.saved_tensors
80
+ # the same constants as used in forward pass.
81
+ floor = -0.043637
82
+ ceil = 1.2
83
+ d = d * ((ceil - floor) / 255.0) + floor
84
+ return y_grad * d
85
+
86
+
87
+ class DoubleSwish(torch.nn.Module):
88
+ def forward(self, x: Tensor) -> Tensor:
89
+ """Return double-swish activation function which is an approximation to Swish(Swish(x)),
90
+ that we approximate closely with x * sigmoid(x-1).
91
+ """
92
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
93
+ return x * torch.sigmoid(x - 1.0)
94
+ return DoubleSwishFunction.apply(x)
95
+
96
+
97
+ class ActivationBalancerFunction(torch.autograd.Function):
98
+ @staticmethod
99
+ def forward(
100
+ ctx,
101
+ x: Tensor,
102
+ scale_factor: Tensor,
103
+ sign_factor: Optional[Tensor],
104
+ channel_dim: int,
105
+ ) -> Tensor:
106
+ if channel_dim < 0:
107
+ channel_dim += x.ndim
108
+ ctx.channel_dim = channel_dim
109
+ xgt0 = x > 0
110
+ if sign_factor is None:
111
+ ctx.save_for_backward(xgt0, scale_factor)
112
+ else:
113
+ ctx.save_for_backward(xgt0, scale_factor, sign_factor)
114
+ return x
115
+
116
+ @staticmethod
117
+ def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
118
+ if len(ctx.saved_tensors) == 3:
119
+ xgt0, scale_factor, sign_factor = ctx.saved_tensors
120
+ for _ in range(ctx.channel_dim, x_grad.ndim - 1):
121
+ scale_factor = scale_factor.unsqueeze(-1)
122
+ sign_factor = sign_factor.unsqueeze(-1)
123
+ factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
124
+ else:
125
+ xgt0, scale_factor = ctx.saved_tensors
126
+ for _ in range(ctx.channel_dim, x_grad.ndim - 1):
127
+ scale_factor = scale_factor.unsqueeze(-1)
128
+ factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
129
+ neg_delta_grad = x_grad.abs() * factor
130
+ return (
131
+ x_grad - neg_delta_grad,
132
+ None,
133
+ None,
134
+ None,
135
+ )
136
+
137
+
138
+ def _compute_scale_factor(
139
+ x: Tensor,
140
+ channel_dim: int,
141
+ min_abs: float,
142
+ max_abs: float,
143
+ gain_factor: float,
144
+ max_factor: float,
145
+ ) -> Tensor:
146
+ if channel_dim < 0:
147
+ channel_dim += x.ndim
148
+ sum_dims = [d for d in range(x.ndim) if d != channel_dim]
149
+ x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
150
+
151
+ if min_abs == 0.0:
152
+ below_threshold = 0.0
153
+ else:
154
+ # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
155
+ # x_abs)_mean , min_abs.
156
+ below_threshold = ((min_abs - x_abs_mean) * (gain_factor / min_abs)).clamp(
157
+ min=0, max=max_factor
158
+ )
159
+
160
+ above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
161
+ min=0, max=max_factor
162
+ )
163
+
164
+ return below_threshold - above_threshold
165
+
166
+
167
+ def _compute_sign_factor(
168
+ x: Tensor,
169
+ channel_dim: int,
170
+ min_positive: float,
171
+ max_positive: float,
172
+ gain_factor: float,
173
+ max_factor: float,
174
+ ) -> Tensor:
175
+ if channel_dim < 0:
176
+ channel_dim += x.ndim
177
+ sum_dims = [d for d in range(x.ndim) if d != channel_dim]
178
+ proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
179
+ if min_positive == 0.0:
180
+ factor1 = 0.0
181
+ else:
182
+ # 0 if proportion_positive >= min_positive, else can be
183
+ # as large as max_factor.
184
+ factor1 = (
185
+ (min_positive - proportion_positive) * (gain_factor / min_positive)
186
+ ).clamp_(min=0, max=max_factor)
187
+
188
+ if max_positive == 1.0:
189
+ factor2 = 0.0
190
+ else:
191
+ # 0 if self.proportion_positive <= max_positive, else can be
192
+ # as large as -max_factor.
193
+ factor2 = (
194
+ (proportion_positive - max_positive) * (gain_factor / (1.0 - max_positive))
195
+ ).clamp_(min=0, max=max_factor)
196
+ sign_factor = factor1 - factor2
197
+ # require min_positive != 0 or max_positive != 1:
198
+ assert not isinstance(sign_factor, float)
199
+ return sign_factor
200
+
201
+
202
+ class ActivationBalancer(torch.nn.Module):
203
+ """
204
+ Modifies the backpropped derivatives of a function to try to encourage, for
205
+ each channel, that it is positive at least a proportion `threshold` of the
206
+ time. It does this by multiplying negative derivative values by up to
207
+ (1+max_factor), and positive derivative values by up to (1-max_factor),
208
+ interpolated from 1 at the threshold to those extremal values when none
209
+ of the inputs are positive.
210
+
211
+ Args:
212
+ num_channels: the number of channels
213
+ channel_dim: the dimension/axis corresponding to the channel, e.g.
214
+ -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
215
+ min_positive: the minimum, per channel, of the proportion of the time
216
+ that (x > 0), below which we start to modify the derivatives.
217
+ max_positive: the maximum, per channel, of the proportion of the time
218
+ that (x > 0), above which we start to modify the derivatives.
219
+ max_factor: the maximum factor by which we modify the derivatives for
220
+ either the sign constraint or the magnitude constraint;
221
+ e.g. with max_factor=0.02, the the derivatives would be multiplied by
222
+ values in the range [0.98..1.02].
223
+ sign_gain_factor: determines the 'gain' with which we increase the
224
+ change in gradient once the constraints on min_positive and max_positive
225
+ are violated.
226
+ scale_gain_factor: determines the 'gain' with which we increase the
227
+ change in gradient once the constraints on min_abs and max_abs
228
+ are violated.
229
+ min_abs: the minimum average-absolute-value difference from the mean
230
+ value per channel, which we allow, before we start to modify
231
+ the derivatives to prevent this.
232
+ max_abs: the maximum average-absolute-value difference from the mean
233
+ value per channel, which we allow, before we start to modify
234
+ the derivatives to prevent this.
235
+ min_prob: determines the minimum probability with which we modify the
236
+ gradients for the {min,max}_positive and {min,max}_abs constraints,
237
+ on each forward(). This is done randomly to prevent all layers
238
+ from doing it at the same time. Early in training we may use
239
+ higher probabilities than this; it will decay to this value.
240
+ """
241
+
242
+ def __init__(
243
+ self,
244
+ num_channels: int,
245
+ channel_dim: int,
246
+ min_positive: float = 0.05,
247
+ max_positive: float = 0.95,
248
+ max_factor: float = 0.04,
249
+ sign_gain_factor: float = 0.01,
250
+ scale_gain_factor: float = 0.02,
251
+ min_abs: float = 0.2,
252
+ max_abs: float = 100.0,
253
+ min_prob: float = 0.1,
254
+ ):
255
+ super(ActivationBalancer, self).__init__()
256
+ self.num_channels = num_channels
257
+ self.channel_dim = channel_dim
258
+ self.min_positive = min_positive
259
+ self.max_positive = max_positive
260
+ self.max_factor = max_factor
261
+ self.min_abs = min_abs
262
+ self.max_abs = max_abs
263
+ self.min_prob = min_prob
264
+ self.sign_gain_factor = sign_gain_factor
265
+ self.scale_gain_factor = scale_gain_factor
266
+
267
+ # count measures how many times the forward() function has been called.
268
+ # We occasionally sync this to a tensor called `count`, that exists to
269
+ # make sure it is synced to disk when we load and save the model.
270
+ self.cpu_count = 0
271
+ self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
272
+
273
+ def forward(self, x: Tensor) -> Tensor:
274
+ if torch.jit.is_scripting() or not x.requires_grad or torch.jit.is_tracing():
275
+ return _no_op(x)
276
+
277
+ count = self.cpu_count
278
+ self.cpu_count += 1
279
+
280
+ if random.random() < 0.01:
281
+ # Occasionally sync self.cpu_count with self.count.
282
+ # count affects the decay of 'prob'. don't do this on every iter,
283
+ # because syncing with the GPU is slow.
284
+ self.cpu_count = max(self.cpu_count, self.count.item())
285
+ self.count.fill_(self.cpu_count)
286
+
287
+ # the prob of doing some work exponentially decreases from 0.5 till it hits
288
+ # a floor at min_prob (==0.1, by default)
289
+ prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
290
+
291
+ if random.random() < prob:
292
+ sign_gain_factor = 0.5
293
+ if self.min_positive != 0.0 or self.max_positive != 1.0:
294
+ sign_factor = _compute_sign_factor(
295
+ x,
296
+ self.channel_dim,
297
+ self.min_positive,
298
+ self.max_positive,
299
+ gain_factor=self.sign_gain_factor / prob,
300
+ max_factor=self.max_factor,
301
+ )
302
+ else:
303
+ sign_factor = None
304
+
305
+ scale_factor = _compute_scale_factor(
306
+ x.detach(),
307
+ self.channel_dim,
308
+ min_abs=self.min_abs,
309
+ max_abs=self.max_abs,
310
+ gain_factor=self.scale_gain_factor / prob,
311
+ max_factor=self.max_factor,
312
+ )
313
+ return ActivationBalancerFunction.apply(
314
+ x,
315
+ scale_factor,
316
+ sign_factor,
317
+ self.channel_dim,
318
+ )
319
+ else:
320
+ return _no_op(x)
321
+
322
+
323
+ def BalancedDoubleSwish(
324
+ d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
325
+ ) -> nn.Sequential:
326
+ """
327
+ ActivationBalancer -> DoubleSwish
328
+ """
329
+ balancer = ActivationBalancer(
330
+ d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
331
+ )
332
+ return nn.Sequential(
333
+ balancer,
334
+ DoubleSwish(),
335
+ )
lipsync-models/GPT_SoVITS/AR/modules/transformer.py ADDED
@@ -0,0 +1,374 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py
2
+ import copy
3
+ import numbers
4
+ from functools import partial
5
+ from typing import Any
6
+ from typing import Callable
7
+ from typing import List
8
+ from typing import Optional
9
+ from typing import Tuple
10
+ from typing import Union
11
+
12
+ import torch
13
+ from AR.modules.activation import MultiheadAttention
14
+ from AR.modules.scaling import BalancedDoubleSwish
15
+ from torch import nn
16
+ from torch import Tensor
17
+ from torch.nn import functional as F
18
+
19
+ _shape_t = Union[int, List[int], torch.Size]
20
+
21
+
22
+ class LayerNorm(nn.Module):
23
+ __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
24
+ normalized_shape: Tuple[int, ...]
25
+ eps: float
26
+ elementwise_affine: bool
27
+
28
+ def __init__(
29
+ self,
30
+ normalized_shape: _shape_t,
31
+ eps: float = 1e-5,
32
+ elementwise_affine: bool = True,
33
+ device=None,
34
+ dtype=None,
35
+ ) -> None:
36
+ factory_kwargs = {"device": device, "dtype": dtype}
37
+ super(LayerNorm, self).__init__()
38
+ if isinstance(normalized_shape, numbers.Integral):
39
+ # mypy error: incompatible types in assignment
40
+ normalized_shape = (normalized_shape,) # type: ignore[assignment]
41
+ self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type]
42
+ self.eps = eps
43
+ self.elementwise_affine = elementwise_affine
44
+ if self.elementwise_affine:
45
+ self.weight = nn.Parameter(
46
+ torch.empty(self.normalized_shape, **factory_kwargs)
47
+ )
48
+ self.bias = nn.Parameter(
49
+ torch.empty(self.normalized_shape, **factory_kwargs)
50
+ )
51
+ else:
52
+ self.register_parameter("weight", None)
53
+ self.register_parameter("bias", None)
54
+
55
+ self.reset_parameters()
56
+
57
+ def reset_parameters(self) -> None:
58
+ if self.elementwise_affine:
59
+ nn.init.ones_(self.weight)
60
+ nn.init.zeros_(self.bias)
61
+
62
+ def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
63
+ if isinstance(input, tuple):
64
+ input, embedding = input
65
+ return (
66
+ F.layer_norm(
67
+ input,
68
+ self.normalized_shape,
69
+ self.weight,
70
+ self.bias,
71
+ self.eps,
72
+ ),
73
+ embedding,
74
+ )
75
+
76
+ assert embedding is None
77
+ return F.layer_norm(
78
+ input, self.normalized_shape, self.weight, self.bias, self.eps
79
+ )
80
+
81
+ def extra_repr(self) -> str:
82
+ return (
83
+ "{normalized_shape}, eps={eps}, "
84
+ "elementwise_affine={elementwise_affine}".format(**self.__dict__)
85
+ )
86
+
87
+
88
+ class IdentityNorm(nn.Module):
89
+ def __init__(
90
+ self,
91
+ d_model: int,
92
+ eps: float = 1e-5,
93
+ device=None,
94
+ dtype=None,
95
+ ) -> None:
96
+ super(IdentityNorm, self).__init__()
97
+
98
+ def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
99
+ if isinstance(input, tuple):
100
+ return input
101
+
102
+ assert embedding is None
103
+ return input
104
+
105
+
106
+ class TransformerEncoder(nn.Module):
107
+ r"""TransformerEncoder is a stack of N encoder layers. Users can build the
108
+ BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
109
+
110
+ Args:
111
+ encoder_layer: an instance of the TransformerEncoderLayer() class (required).
112
+ num_layers: the number of sub-encoder-layers in the encoder (required).
113
+ norm: the layer normalization component (optional).
114
+ enable_nested_tensor: if True, input will automatically convert to nested tensor
115
+ (and convert back on output). This will improve the overall performance of
116
+ TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
117
+
118
+ Examples::
119
+ >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
120
+ >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
121
+ >>> src = torch.rand(10, 32, 512)
122
+ >>> out = transformer_encoder(src)
123
+ """
124
+ __constants__ = ["norm"]
125
+
126
+ def __init__(self, encoder_layer, num_layers, norm=None):
127
+ super(TransformerEncoder, self).__init__()
128
+ self.layers = _get_clones(encoder_layer, num_layers)
129
+ self.num_layers = num_layers
130
+ self.norm = norm
131
+
132
+ # def forward(
133
+ # self,
134
+ # src: Tensor,
135
+ # mask: Optional[Tensor] = None,
136
+ # src_key_padding_mask: Optional[Tensor] = None,
137
+ # return_layer_states: bool = False,
138
+ # cache=None,
139
+ # ) -> Tensor:
140
+ # r"""Pass the input through the encoder layers in turn.
141
+
142
+ # Args:
143
+ # src: the sequence to the encoder (required).
144
+ # mask: the mask for the src sequence (optional).
145
+ # src_key_padding_mask: the mask for the src keys per batch (optional).
146
+ # return_layer_states: return layers' state (optional).
147
+
148
+ # Shape:
149
+ # see the docs in Transformer class.
150
+ # """
151
+ # if return_layer_states:
152
+ # layer_states = [] # layers' output
153
+ # output = src
154
+ # for mod in self.layers:
155
+ # output, _ = mod(
156
+ # output,
157
+ # src_mask=mask,
158
+ # src_key_padding_mask=src_key_padding_mask,
159
+ # cache=cache,
160
+ # )
161
+ # layer_states.append(output[0])
162
+
163
+ # if self.norm is not None:
164
+ # output = self.norm(output)
165
+
166
+ # return layer_states, output
167
+
168
+ # output = src
169
+ # for mod in self.layers:
170
+ # output = mod(
171
+ # output,
172
+ # src_mask=mask,
173
+ # src_key_padding_mask=src_key_padding_mask,
174
+ # cache=cache,
175
+ # )
176
+
177
+ # if self.norm is not None:
178
+ # output = self.norm(output)
179
+
180
+ # return output
181
+
182
+
183
+ class TransformerEncoderLayer(nn.Module):
184
+ __constants__ = ["batch_first", "norm_first"]
185
+
186
+ def __init__(
187
+ self,
188
+ d_model: int,
189
+ nhead: int,
190
+ dim_feedforward: int = 2048,
191
+ dropout: float = 0.1,
192
+ activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
193
+ batch_first: bool = False,
194
+ norm_first: bool = False,
195
+ device=None,
196
+ dtype=None,
197
+ linear1_self_attention_cls: nn.Module = nn.Linear,
198
+ linear2_self_attention_cls: nn.Module = nn.Linear,
199
+ linear1_feedforward_cls: nn.Module = nn.Linear,
200
+ linear2_feedforward_cls: nn.Module = nn.Linear,
201
+ layer_norm_cls: nn.Module = LayerNorm,
202
+ layer_norm_eps: float = 1e-5,
203
+ adaptive_layer_norm=False,
204
+ ) -> None:
205
+ factory_kwargs = {"device": device, "dtype": dtype}
206
+ super(TransformerEncoderLayer, self).__init__()
207
+ # print(233333333333,d_model,nhead)
208
+ # import os
209
+ # os._exit(2333333)
210
+ self.self_attn = MultiheadAttention(
211
+ d_model, # 512 16
212
+ nhead,
213
+ dropout=dropout,
214
+ batch_first=batch_first,
215
+ linear1_cls=linear1_self_attention_cls,
216
+ linear2_cls=linear2_self_attention_cls,
217
+ **factory_kwargs,
218
+ )
219
+
220
+ # Implementation of Feedforward model
221
+ self.linear1 = linear1_feedforward_cls(
222
+ d_model, dim_feedforward, **factory_kwargs
223
+ )
224
+ self.dropout = nn.Dropout(dropout)
225
+ self.linear2 = linear2_feedforward_cls(
226
+ dim_feedforward, d_model, **factory_kwargs
227
+ )
228
+
229
+ self.norm_first = norm_first
230
+ self.dropout1 = nn.Dropout(dropout)
231
+ self.dropout2 = nn.Dropout(dropout)
232
+
233
+ # Legacy string support for activation function.
234
+ if isinstance(activation, str):
235
+ activation = _get_activation_fn(activation)
236
+ elif isinstance(activation, partial):
237
+ activation = activation(d_model)
238
+ elif activation == BalancedDoubleSwish:
239
+ activation = BalancedDoubleSwish(d_model)
240
+
241
+ # # We can't test self.activation in forward() in TorchScript,
242
+ # # so stash some information about it instead.
243
+ # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
244
+ # self.activation_relu_or_gelu = 1
245
+ # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
246
+ # self.activation_relu_or_gelu = 2
247
+ # else:
248
+ # self.activation_relu_or_gelu = 0
249
+ self.activation = activation
250
+
251
+ norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
252
+ if layer_norm_cls == IdentityNorm:
253
+ norm2 = BalancedBasicNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
254
+ else:
255
+ norm2 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
256
+
257
+ if adaptive_layer_norm:
258
+ self.norm1 = AdaptiveLayerNorm(d_model, norm1)
259
+ self.norm2 = AdaptiveLayerNorm(d_model, norm2)
260
+ else:
261
+ self.norm1 = norm1
262
+ self.norm2 = norm2
263
+
264
+ def __setstate__(self, state):
265
+ super(TransformerEncoderLayer, self).__setstate__(state)
266
+ if not hasattr(self, "activation"):
267
+ self.activation = F.relu
268
+
269
+ def forward(
270
+ self,
271
+ src: Tensor,
272
+ src_mask: Optional[Tensor] = None,
273
+ src_key_padding_mask: Optional[Tensor] = None,
274
+ cache=None,
275
+ ) -> Tuple[Tensor, Optional[Tensor]]:
276
+ r"""Pass the input through the encoder layer.
277
+
278
+ Args:
279
+ src: the sequence to the encoder layer (required).
280
+ src_mask: the mask for the src sequence (optional).
281
+ src_key_padding_mask: the mask for the src keys per batch (optional).
282
+
283
+ Shape:
284
+ see the docs in Transformer class.
285
+ """
286
+ x, stage_embedding = src, None
287
+ if isinstance(src, tuple):
288
+ x, stage_embedding = src
289
+
290
+ if src_key_padding_mask is not None:
291
+ _skpm_dtype = src_key_padding_mask.dtype
292
+ if _skpm_dtype != torch.bool and not torch.is_floating_point(
293
+ src_key_padding_mask
294
+ ):
295
+ raise AssertionError(
296
+ "only bool and floating types of key_padding_mask are supported"
297
+ )
298
+
299
+ if self.norm_first:
300
+ x = x + self._sa_block(
301
+ self.norm1(x, stage_embedding),
302
+ src_mask,
303
+ src_key_padding_mask,
304
+ cache=cache,
305
+ )
306
+ x = x + self._ff_block(self.norm2(x, stage_embedding))
307
+ else:
308
+ x = self.norm1(
309
+ x + self._sa_block(x, src_mask, src_key_padding_mask, cache=cache),
310
+ stage_embedding,
311
+ )
312
+ x = self.norm2(x + self._ff_block(x), stage_embedding)
313
+
314
+ return (x, stage_embedding)
315
+
316
+ # self-attention block
317
+ def _sa_block(
318
+ self,
319
+ x: Tensor,
320
+ attn_mask: Optional[Tensor],
321
+ key_padding_mask: Optional[Tensor],
322
+ cache=None,
323
+ ) -> Tensor:
324
+ # print(x.shape,attn_mask.shape,key_padding_mask)
325
+ # torch.Size([1, 188, 512]) torch.Size([188, 188]) None
326
+ # import os
327
+ # os._exit(23333)
328
+ x = self.self_attn(
329
+ x,
330
+ x,
331
+ x,
332
+ attn_mask=attn_mask,
333
+ key_padding_mask=key_padding_mask,
334
+ need_weights=False,
335
+ cache=cache,
336
+ )[0]
337
+ return self.dropout1(x)
338
+
339
+ # feed forward block
340
+ def _ff_block(self, x: Tensor) -> Tensor:
341
+ x = self.linear2(self.dropout(self.activation(self.linear1(x))))
342
+ return self.dropout2(x)
343
+
344
+
345
+ class AdaptiveLayerNorm(nn.Module):
346
+ r"""Adaptive Layer Normalization"""
347
+
348
+ def __init__(self, d_model, norm) -> None:
349
+ super(AdaptiveLayerNorm, self).__init__()
350
+ self.project_layer = nn.Linear(d_model, 2 * d_model)
351
+ self.norm = norm
352
+ self.d_model = d_model
353
+ self.eps = self.norm.eps
354
+
355
+ def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
356
+ if isinstance(input, tuple):
357
+ input, embedding = input
358
+ weight, bias = torch.split(
359
+ self.project_layer(embedding),
360
+ split_size_or_sections=self.d_model,
361
+ dim=-1,
362
+ )
363
+ return (weight * self.norm(input) + bias, embedding)
364
+
365
+ weight, bias = torch.split(
366
+ self.project_layer(embedding),
367
+ split_size_or_sections=self.d_model,
368
+ dim=-1,
369
+ )
370
+ return weight * self.norm(input) + bias
371
+
372
+
373
+ def _get_clones(module, N):
374
+ return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
lipsync-models/GPT_SoVITS/AR/text_processing/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/AR/text_processing/phonemizer.py ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/text_processing/phonemizer.py
2
+ import itertools
3
+ import re
4
+ from typing import Dict
5
+ from typing import List
6
+
7
+ import regex
8
+ from gruut import sentences
9
+ from gruut.const import Sentence
10
+ from gruut.const import Word
11
+ from AR.text_processing.symbols import SYMBOL_TO_ID
12
+
13
+
14
+ class GruutPhonemizer:
15
+ def __init__(self, language: str):
16
+ self._phonemizer = sentences
17
+ self.lang = language
18
+ self.symbol_to_id = SYMBOL_TO_ID
19
+ self._special_cases_dict: Dict[str] = {
20
+ r"\.\.\.": "... ",
21
+ ";": "; ",
22
+ ":": ": ",
23
+ ",": ", ",
24
+ r"\.": ". ",
25
+ "!": "! ",
26
+ r"\?": "? ",
27
+ "—": "—",
28
+ "…": "… ",
29
+ "«": "«",
30
+ "»": "»",
31
+ }
32
+ self._punctuation_regexp: str = (
33
+ rf"([{''.join(self._special_cases_dict.keys())}])"
34
+ )
35
+
36
+ def _normalize_punctuation(self, text: str) -> str:
37
+ text = regex.sub(rf"\pZ+{self._punctuation_regexp}", r"\1", text)
38
+ text = regex.sub(rf"{self._punctuation_regexp}(\pL)", r"\1 \2", text)
39
+ text = regex.sub(r"\pZ+", r" ", text)
40
+ return text.strip()
41
+
42
+ def _convert_punctuation(self, word: Word) -> str:
43
+ if not word.phonemes:
44
+ return ""
45
+ if word.phonemes[0] in ["‖", "|"]:
46
+ return word.text.strip()
47
+
48
+ phonemes = "".join(word.phonemes)
49
+ # remove modifier characters ˈˌː with regex
50
+ phonemes = re.sub(r"[ˈˌː͡]", "", phonemes)
51
+ return phonemes.strip()
52
+
53
+ def phonemize(self, text: str, espeak: bool = False) -> str:
54
+ text_to_phonemize: str = self._normalize_punctuation(text)
55
+ sents: List[Sentence] = [
56
+ sent
57
+ for sent in self._phonemizer(text_to_phonemize, lang="en-us", espeak=espeak)
58
+ ]
59
+ words: List[str] = [
60
+ self._convert_punctuation(word) for word in itertools.chain(*sents)
61
+ ]
62
+ return " ".join(words)
63
+
64
+ def transform(self, phonemes):
65
+ # convert phonemes to ids
66
+ # dictionary is in symbols.py
67
+ return [self.symbol_to_id[p] for p in phonemes if p in self.symbol_to_id.keys()]
68
+
69
+
70
+ if __name__ == "__main__":
71
+ phonemizer = GruutPhonemizer("en-us")
72
+ # text -> IPA
73
+ phonemes = phonemizer.phonemize("Hello, wor-ld ?")
74
+ print("phonemes:", phonemes)
75
+ print("len(phonemes):", len(phonemes))
76
+ phoneme_ids = phonemizer.transform(phonemes)
77
+ print("phoneme_ids:", phoneme_ids)
78
+ print("len(phoneme_ids):", len(phoneme_ids))
lipsync-models/GPT_SoVITS/AR/text_processing/symbols.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ # modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/text_processing/symbols.py
2
+ PAD = "_"
3
+ PUNCTUATION = ';:,.!?¡¿—…"«»“” '
4
+ LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
5
+ IPA_LETTERS = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
6
+ SYMBOLS = [PAD] + list(PUNCTUATION) + list(LETTERS) + list(IPA_LETTERS)
7
+ SPACE_ID = SYMBOLS.index(" ")
8
+ SYMBOL_TO_ID = {s: i for i, s in enumerate(SYMBOLS)}
9
+ ID_TO_SYMBOL = {i: s for i, s in enumerate(SYMBOLS)}
lipsync-models/GPT_SoVITS/AR/utils/__init__.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+
3
+
4
+ def str2bool(str):
5
+ return True if str.lower() == 'true' else False
6
+
7
+
8
+ def get_newest_ckpt(string_list):
9
+ # 定义一个正则表达式模式,用于匹配字符串中的数字
10
+ pattern = r'epoch=(\d+)-step=(\d+)\.ckpt'
11
+
12
+ # 使用正则表达式提取每个字符串中的数字信息,并创建一个包含元组的列表
13
+ extracted_info = []
14
+ for string in string_list:
15
+ match = re.match(pattern, string)
16
+ if match:
17
+ epoch = int(match.group(1))
18
+ step = int(match.group(2))
19
+ extracted_info.append((epoch, step, string))
20
+ # 按照 epoch 后面的数字和 step 后面的数字进行排序
21
+ sorted_info = sorted(
22
+ extracted_info, key=lambda x: (x[0], x[1]), reverse=True)
23
+ # 获取最新的 ckpt 文件名
24
+ newest_ckpt = sorted_info[0][2]
25
+ return newest_ckpt
26
+
27
+
28
+ # 文本存在且不为空时 return True
29
+ def check_txt_file(file_path):
30
+ try:
31
+ with open(file_path, 'r') as file:
32
+ text = file.readline().strip()
33
+ assert text.strip() != ''
34
+ return text
35
+ except Exception:
36
+ return False
37
+ return False
lipsync-models/GPT_SoVITS/AR/utils/initialize.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Initialize modules for espnet2 neural networks."""
3
+ import torch
4
+ from typeguard import check_argument_types
5
+
6
+
7
+ def initialize(model: torch.nn.Module, init: str):
8
+ """Initialize weights of a neural network module.
9
+
10
+ Parameters are initialized using the given method or distribution.
11
+
12
+ Custom initialization routines can be implemented into submodules
13
+ as function `espnet_initialization_fn` within the custom module.
14
+
15
+ Args:
16
+ model: Target.
17
+ init: Method of initialization.
18
+ """
19
+ assert check_argument_types()
20
+ print("init with", init)
21
+
22
+ # weight init
23
+ for p in model.parameters():
24
+ if p.dim() > 1:
25
+ if init == "xavier_uniform":
26
+ torch.nn.init.xavier_uniform_(p.data)
27
+ elif init == "xavier_normal":
28
+ torch.nn.init.xavier_normal_(p.data)
29
+ elif init == "kaiming_uniform":
30
+ torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu")
31
+ elif init == "kaiming_normal":
32
+ torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu")
33
+ else:
34
+ raise ValueError("Unknown initialization: " + init)
35
+ # bias init
36
+ for name, p in model.named_parameters():
37
+ if ".bias" in name and p.dim() == 1:
38
+ p.data.zero_()
lipsync-models/GPT_SoVITS/AR/utils/io.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+
3
+ import torch
4
+ import yaml
5
+
6
+
7
+ def load_yaml_config(path):
8
+ with open(path) as f:
9
+ config = yaml.full_load(f)
10
+ return config
11
+
12
+
13
+ def save_config_to_yaml(config, path):
14
+ assert path.endswith(".yaml")
15
+ with open(path, "w") as f:
16
+ f.write(yaml.dump(config))
17
+ f.close()
18
+
19
+
20
+ def write_args(args, path):
21
+ args_dict = dict(
22
+ (name, getattr(args, name)) for name in dir(args) if not name.startswith("_")
23
+ )
24
+ with open(path, "a") as args_file:
25
+ args_file.write("==> torch version: {}\n".format(torch.__version__))
26
+ args_file.write(
27
+ "==> cudnn version: {}\n".format(torch.backends.cudnn.version())
28
+ )
29
+ args_file.write("==> Cmd:\n")
30
+ args_file.write(str(sys.argv))
31
+ args_file.write("\n==> args:\n")
32
+ for k, v in sorted(args_dict.items()):
33
+ args_file.write(" %s: %s\n" % (str(k), str(v)))
34
+ args_file.close()
lipsync-models/GPT_SoVITS/configs/s1.yaml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ train:
2
+ seed: 1234
3
+ epochs: 300
4
+ batch_size: 8
5
+ gradient_accumulation: 4
6
+ save_every_n_epoch: 1
7
+ precision: 16
8
+ gradient_clip: 1.0
9
+ optimizer:
10
+ lr: 0.01
11
+ lr_init: 0.00001
12
+ lr_end: 0.0001
13
+ warmup_steps: 2000
14
+ decay_steps: 40000
15
+ data:
16
+ max_eval_sample: 8
17
+ max_sec: 54
18
+ num_workers: 1
19
+ pad_val: 1024 # same with EOS in model
20
+ model:
21
+ vocab_size: 1025
22
+ phoneme_vocab_size: 512
23
+ embedding_dim: 512
24
+ hidden_dim: 512
25
+ head: 16
26
+ linear_units: 2048
27
+ n_layer: 12
28
+ dropout: 0
29
+ EOS: 1024
30
+ inference:
31
+ top_k: 5
lipsync-models/GPT_SoVITS/configs/s1big.yaml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ train:
2
+ seed: 1234
3
+ epochs: 300
4
+ batch_size: 8
5
+ gradient_accumulation: 4
6
+ save_every_n_epoch: 1
7
+ precision: 16-mixed
8
+ gradient_clip: 1.0
9
+ optimizer:
10
+ lr: 0.01
11
+ lr_init: 0.00001
12
+ lr_end: 0.0001
13
+ warmup_steps: 2000
14
+ decay_steps: 40000
15
+ data:
16
+ max_eval_sample: 8
17
+ max_sec: 54
18
+ num_workers: 1
19
+ pad_val: 1024 # same with EOS in model
20
+ model:
21
+ vocab_size: 1025
22
+ phoneme_vocab_size: 512
23
+ embedding_dim: 1024
24
+ hidden_dim: 1024
25
+ head: 16
26
+ linear_units: 2048
27
+ n_layer: 16
28
+ dropout: 0
29
+ EOS: 1024
30
+ inference:
31
+ top_k: 5
lipsync-models/GPT_SoVITS/configs/s1big2.yaml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ train:
2
+ seed: 1234
3
+ epochs: 300
4
+ batch_size: 12
5
+ gradient_accumulation: 4
6
+ save_every_n_epoch: 1
7
+ precision: 16-mixed
8
+ gradient_clip: 1.0
9
+ optimizer:
10
+ lr: 0.01
11
+ lr_init: 0.00001
12
+ lr_end: 0.0001
13
+ warmup_steps: 2000
14
+ decay_steps: 40000
15
+ data:
16
+ max_eval_sample: 8
17
+ max_sec: 54
18
+ num_workers: 1
19
+ pad_val: 1024 # same with EOS in model
20
+ model:
21
+ vocab_size: 1025
22
+ phoneme_vocab_size: 512
23
+ embedding_dim: 1024
24
+ hidden_dim: 1024
25
+ head: 16
26
+ linear_units: 2048
27
+ n_layer: 6
28
+ dropout: 0
29
+ EOS: 1024
30
+ inference:
31
+ top_k: 5
lipsync-models/GPT_SoVITS/configs/s1longer.yaml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ train:
2
+ seed: 1234
3
+ epochs: 20
4
+ batch_size: 8
5
+ save_every_n_epoch: 1
6
+ precision: 16-mixed
7
+ gradient_clip: 1.0
8
+ optimizer:
9
+ lr: 0.01
10
+ lr_init: 0.00001
11
+ lr_end: 0.0001
12
+ warmup_steps: 2000
13
+ decay_steps: 40000
14
+ data:
15
+ max_eval_sample: 8
16
+ max_sec: 54
17
+ num_workers: 4
18
+ pad_val: 1024 # same with EOS in model
19
+ model:
20
+ vocab_size: 1025
21
+ phoneme_vocab_size: 512
22
+ embedding_dim: 512
23
+ hidden_dim: 512
24
+ head: 16
25
+ linear_units: 2048
26
+ n_layer: 24
27
+ dropout: 0
28
+ EOS: 1024
29
+ random_bert: 0
30
+ inference:
31
+ top_k: 5
lipsync-models/GPT_SoVITS/configs/s1mq.yaml ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ train:
2
+ seed: 1234
3
+ epochs: 100
4
+ batch_size: 6
5
+ gradient_accumulation: 4
6
+ save_every_n_epoch: 1
7
+ precision: 32
8
+ gradient_clip: 1.0
9
+ optimizer:
10
+ lr: 0.01
11
+ lr_init: 0.00001
12
+ lr_end: 0.0001
13
+ warmup_steps: 2000
14
+ decay_steps: 40000
15
+ data:
16
+ max_eval_sample: 8
17
+ max_sec: 40
18
+ num_workers: 1
19
+ pad_val: 1024 # same with EOS in model
20
+ model:
21
+ saving_path: "ckpt/"
22
+ resume_checkpoint: null
23
+ vocoder_config_path: "quantizer/new_ckpt/config.json"
24
+ vocoder_ckpt_path: "quantizer/new_ckpt/g_00600000"
25
+ datadir: "/home/liweiche/GigaSpeech/wavs"
26
+ metapath: "/home/liweiche/GigaSpeech/train2.json"
27
+ val_metapath: "/home/liweiche/GigaSpeech/dev2.json"
28
+ sampledir: "logs/"
29
+ pretrained_path: null
30
+ lr: 0.0001
31
+ batch_size: 200.0
32
+ train_bucket_size: 8192
33
+ training_step: 800000
34
+ optim_flat_percent: 0.0
35
+ warmup_step: 50
36
+ adam_beta1: 0.9
37
+ adam_beta2: 0.98
38
+ ffd_size: 3072
39
+ hidden_size: 768
40
+ enc_nlayers: 6
41
+ dec_nlayers: 6
42
+ nheads: 12
43
+ ar_layer: 4
44
+ ar_ffd_size: 1024
45
+ ar_hidden_size: 256
46
+ ar_nheads: 4
47
+ aligner_softmax_temp: 1.0
48
+ layer_norm_eps: 0.00001
49
+ speaker_embed_dropout: 0.05
50
+ label_smoothing: 0.0
51
+ val_check_interval: 5000
52
+ check_val_every_n_epoch: 1
53
+ precision: "fp16"
54
+ nworkers: 16
55
+ distributed: true
56
+ accelerator: "ddp"
57
+ version: null
58
+ accumulate_grad_batches: 1
59
+ use_repetition_token: true
60
+ use_repetition_gating: false
61
+ repetition_penalty: 1.0
62
+ sampling_temperature: 1.0
63
+ top_k: -1
64
+ min_top_k: 3
65
+ top_p: 0.8
66
+ sample_num: 4
67
+ length_penalty_max_length: 15000
68
+ length_penalty_max_prob: 0.95
69
+ max_input_length: 2048
70
+ max_output_length: 2000
71
+ sample_rate: 16000
72
+ n_codes: 1024
73
+ n_cluster_groups: 1
74
+ phone_context_window: 4
75
+ phoneset_size: 1000
76
+ inference:
77
+ top_k: 5
lipsync-models/GPT_SoVITS/configs/s2.json ADDED
@@ -0,0 +1,90 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "train": {
3
+ "log_interval": 100,
4
+ "eval_interval": 500,
5
+ "seed": 1234,
6
+ "epochs": 100,
7
+ "learning_rate": 0.0001,
8
+ "betas": [
9
+ 0.8,
10
+ 0.99
11
+ ],
12
+ "eps": 1e-09,
13
+ "batch_size": 32,
14
+ "fp16_run": true,
15
+ "lr_decay": 0.999875,
16
+ "segment_size": 20480,
17
+ "init_lr_ratio": 1,
18
+ "warmup_epochs": 0,
19
+ "c_mel": 45,
20
+ "c_kl": 1.0,
21
+ "text_low_lr_rate": 0.4
22
+ },
23
+ "data": {
24
+ "max_wav_value": 32768.0,
25
+ "sampling_rate": 32000,
26
+ "filter_length": 2048,
27
+ "hop_length": 640,
28
+ "win_length": 2048,
29
+ "n_mel_channels": 128,
30
+ "mel_fmin": 0.0,
31
+ "mel_fmax": null,
32
+ "add_blank": true,
33
+ "n_speakers": 300,
34
+ "cleaned_text": true
35
+ },
36
+ "model": {
37
+ "inter_channels": 192,
38
+ "hidden_channels": 192,
39
+ "filter_channels": 768,
40
+ "n_heads": 2,
41
+ "n_layers": 6,
42
+ "kernel_size": 3,
43
+ "p_dropout": 0.1,
44
+ "resblock": "1",
45
+ "resblock_kernel_sizes": [
46
+ 3,
47
+ 7,
48
+ 11
49
+ ],
50
+ "resblock_dilation_sizes": [
51
+ [
52
+ 1,
53
+ 3,
54
+ 5
55
+ ],
56
+ [
57
+ 1,
58
+ 3,
59
+ 5
60
+ ],
61
+ [
62
+ 1,
63
+ 3,
64
+ 5
65
+ ]
66
+ ],
67
+ "upsample_rates": [
68
+ 10,
69
+ 8,
70
+ 2,
71
+ 2,
72
+ 2
73
+ ],
74
+ "upsample_initial_channel": 512,
75
+ "upsample_kernel_sizes": [
76
+ 16,
77
+ 16,
78
+ 8,
79
+ 2,
80
+ 2
81
+ ],
82
+ "n_layers_q": 3,
83
+ "use_spectral_norm": false,
84
+ "gin_channels": 512,
85
+ "semantic_frame_rate": "25hz",
86
+ "freeze_quantizer": true
87
+ },
88
+ "s2_ckpt_dir": "logs/s2/big2k1",
89
+ "content_module": "cnhubert"
90
+ }
lipsync-models/GPT_SoVITS/configs/train.yaml ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gpu:
2
+ n_card: 1
3
+ n_process_per_card: 2
4
+ io:
5
+ text_path: D:\RVC1006\GPT-SoVITS\GPT_SoVITS
6
+ save_every_n_epoch: 1
7
+ precision: 16-mixed
8
+ gradient_clip: 1.0
9
+ optimizer:
10
+ lr: 0.01
11
+ lr_init: 0.00001
12
+ lr_end: 0.0001
13
+ warmup_steps: 2000
14
+ decay_steps: 40000
15
+ data:
16
+ max_eval_sample: 8
17
+ max_sec: 54
18
+ num_workers: 1
19
+ pad_val: 1024 # same with EOS in model
20
+ model:
21
+ vocab_size: 1025
22
+ phoneme_vocab_size: 512
23
+ embedding_dim: 512
24
+ hidden_dim: 512
25
+ head: 16
26
+ linear_units: 2048
27
+ n_layer: 24
28
+ dropout: 0
29
+ EOS: 1024
30
+ random_bert: 0
31
+ inference:
32
+ top_k: 5
lipsync-models/GPT_SoVITS/feature_extractor/__init__.py ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ from . import cnhubert, whisper_enc
2
+
3
+ content_module_map = {
4
+ 'cnhubert': cnhubert,
5
+ 'whisper': whisper_enc
6
+ }
lipsync-models/GPT_SoVITS/feature_extractor/cnhubert.py ADDED
@@ -0,0 +1,104 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+
3
+ import librosa
4
+ import torch
5
+ import torch.nn.functional as F
6
+ import soundfile as sf
7
+ import logging
8
+
9
+ logging.getLogger("numba").setLevel(logging.WARNING)
10
+
11
+ from transformers import (
12
+ Wav2Vec2FeatureExtractor,
13
+ HubertModel,
14
+ )
15
+
16
+ import utils
17
+ import torch.nn as nn
18
+
19
+ cnhubert_base_path = None
20
+
21
+
22
+ class CNHubert(nn.Module):
23
+ def __init__(self):
24
+ super().__init__()
25
+ self.model = HubertModel.from_pretrained(cnhubert_base_path)
26
+ self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
27
+ cnhubert_base_path
28
+ )
29
+
30
+ def forward(self, x):
31
+ input_values = self.feature_extractor(
32
+ x, return_tensors="pt", sampling_rate=16000
33
+ ).input_values.to(x.device)
34
+ feats = self.model(input_values)["last_hidden_state"]
35
+ return feats
36
+
37
+
38
+ # class CNHubertLarge(nn.Module):
39
+ # def __init__(self):
40
+ # super().__init__()
41
+ # self.model = HubertModel.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-hubert-large")
42
+ # self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-hubert-large")
43
+ # def forward(self, x):
44
+ # input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
45
+ # feats = self.model(input_values)["last_hidden_state"]
46
+ # return feats
47
+ #
48
+ # class CVec(nn.Module):
49
+ # def __init__(self):
50
+ # super().__init__()
51
+ # self.model = HubertModel.from_pretrained("/data/docker/liujing04/vc-webui-big/hubert_base")
52
+ # self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/vc-webui-big/hubert_base")
53
+ # def forward(self, x):
54
+ # input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
55
+ # feats = self.model(input_values)["last_hidden_state"]
56
+ # return feats
57
+ #
58
+ # class cnw2v2base(nn.Module):
59
+ # def __init__(self):
60
+ # super().__init__()
61
+ # self.model = Wav2Vec2Model.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-wav2vec2-base")
62
+ # self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-wav2vec2-base")
63
+ # def forward(self, x):
64
+ # input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
65
+ # feats = self.model(input_values)["last_hidden_state"]
66
+ # return feats
67
+
68
+
69
+ def get_model():
70
+ model = CNHubert()
71
+ model.eval()
72
+ return model
73
+
74
+
75
+ # def get_large_model():
76
+ # model = CNHubertLarge()
77
+ # model.eval()
78
+ # return model
79
+ #
80
+ # def get_model_cvec():
81
+ # model = CVec()
82
+ # model.eval()
83
+ # return model
84
+ #
85
+ # def get_model_cnw2v2base():
86
+ # model = cnw2v2base()
87
+ # model.eval()
88
+ # return model
89
+
90
+
91
+ def get_content(hmodel, wav_16k_tensor):
92
+ with torch.no_grad():
93
+ feats = hmodel(wav_16k_tensor)
94
+ return feats.transpose(1, 2)
95
+
96
+
97
+ if __name__ == "__main__":
98
+ model = get_model()
99
+ src_path = "/Users/Shared/原音频2.wav"
100
+ wav_16k_tensor = utils.load_wav_to_torch_and_resample(src_path, 16000)
101
+ model = model
102
+ wav_16k_tensor = wav_16k_tensor
103
+ feats = get_content(model, wav_16k_tensor)
104
+ print(feats.shape)
lipsync-models/GPT_SoVITS/feature_extractor/whisper_enc.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+
4
+ def get_model():
5
+ import whisper
6
+
7
+ model = whisper.load_model("small", device="cpu")
8
+
9
+ return model.encoder
10
+
11
+
12
+ def get_content(model=None, wav_16k_tensor=None):
13
+ from whisper import log_mel_spectrogram, pad_or_trim
14
+
15
+ dev = next(model.parameters()).device
16
+ mel = log_mel_spectrogram(wav_16k_tensor).to(dev)[:, :3000]
17
+ # if torch.cuda.is_available():
18
+ # mel = mel.to(torch.float16)
19
+ feature_len = mel.shape[-1] // 2
20
+ assert mel.shape[-1] < 3000, "输入音频过长,只允许输入30以内音频"
21
+ with torch.no_grad():
22
+ feature = model(pad_or_trim(mel, 3000).unsqueeze(0))[
23
+ :1, :feature_len, :
24
+ ].transpose(1, 2)
25
+ return feature
lipsync-models/GPT_SoVITS/gpt_sovits.py ADDED
@@ -0,0 +1,445 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import string
3
+
4
+ import emoji
5
+ import langid
6
+ import LangSegment
7
+ import librosa
8
+ import numpy as np
9
+ import preprocess
10
+ import soundfile as sf
11
+ import torch
12
+ from AR.models.t2s_lightning_module import Text2SemanticLightningModule
13
+ from feature_extractor import cnhubert
14
+ from module.models import SynthesizerTrn
15
+ from safetensors import safe_open
16
+ from safetensors.torch import save_file
17
+ from transformers import AutoModelForMaskedLM, AutoTokenizer
18
+
19
+
20
+ class TTSConfig(object):
21
+ default_sovits_path = "GPT_SoVITS/pretrained_models/s2G488k.pth"
22
+ default_gpt_path = (
23
+ "GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"
24
+ )
25
+ cnhubert_path = "GPT_SoVITS/pretrained_models/chinese-hubert-base"
26
+ bert_path = "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large"
27
+
28
+ def __init__(self, is_half=True, device="cuda", sovits_path=None, gpt_path=None):
29
+ self.is_half = is_half
30
+ self.device = device
31
+ self.torch_dtype = torch.float16 if is_half == True else torch.float32
32
+ self.numpy_dtype = np.float16 if is_half == True else np.float32
33
+ self.sovits_path = (
34
+ TTSConfig.default_sovits_path if sovits_path is None else sovits_path
35
+ )
36
+ self.gpt_path = TTSConfig.default_gpt_path if gpt_path is None else gpt_path
37
+
38
+
39
+ class DictToAttrRecursive(dict):
40
+ def __init__(self, input_dict):
41
+ super().__init__(input_dict)
42
+ for key, value in input_dict.items():
43
+ if isinstance(value, dict):
44
+ value = DictToAttrRecursive(value)
45
+ self[key] = value
46
+ setattr(self, key, value)
47
+
48
+ def __getattr__(self, item):
49
+ try:
50
+ return self[item]
51
+ except KeyError:
52
+ raise AttributeError(f"Attribute {item} not found")
53
+
54
+ def __setattr__(self, key, value):
55
+ if isinstance(value, dict):
56
+ value = DictToAttrRecursive(value)
57
+ super(DictToAttrRecursive, self).__setitem__(key, value)
58
+ super().__setattr__(key, value)
59
+
60
+ def __delattr__(self, item):
61
+ try:
62
+ del self[item]
63
+ except KeyError:
64
+ raise AttributeError(f"Attribute {item} not found")
65
+
66
+
67
+ class Speaker(object):
68
+ def __init__(
69
+ self,
70
+ refer,
71
+ prompt_semantic,
72
+ phoneme,
73
+ bert,
74
+ audio_path: str = None,
75
+ speaker_name: str = None,
76
+ ):
77
+ self.refer = refer
78
+ self.prompt_semantic = prompt_semantic
79
+ self.phoneme = phoneme
80
+ self.bert = bert
81
+ self.audio_path = audio_path
82
+ self.speaker_name = speaker_name
83
+
84
+ def save(self, path):
85
+ torch_phoneme = torch.LongTensor(self.phoneme)
86
+
87
+ tensors = {
88
+ "refer": self.refer,
89
+ "prompt_semantic": self.prompt_semantic,
90
+ "phoneme": torch_phoneme,
91
+ "bert": self.bert,
92
+ }
93
+ save_file(tensors, path)
94
+
95
+ @staticmethod
96
+ def load(path, device="cpu"):
97
+ with safe_open(path, framework="pt", device="cpu") as tensors:
98
+ return Speaker(
99
+ tensors.get_tensor("refer").to(device),
100
+ tensors.get_tensor("prompt_semantic").to(device),
101
+ tensors.get_tensor("phoneme").tolist(),
102
+ tensors.get_tensor("bert").to(device),
103
+ )
104
+
105
+
106
+ class TTS(object):
107
+ def __init__(self, config: TTSConfig):
108
+ self.config = config
109
+ self.load_models()
110
+
111
+ def load_models(self):
112
+ cnhubert.cnhubert_base_path = self.config.cnhubert_path
113
+ tokenizer = AutoTokenizer.from_pretrained(self.config.bert_path)
114
+ bert_model = AutoModelForMaskedLM.from_pretrained(self.config.bert_path)
115
+ ssl_model = cnhubert.get_model()
116
+
117
+ dict_s2 = torch.load(self.config.sovits_path, map_location="cpu")
118
+ hps = dict_s2["config"]
119
+ hps = DictToAttrRecursive(hps)
120
+ hps.model.semantic_frame_rate = "25hz"
121
+
122
+ dict_s1 = torch.load(self.config.gpt_path, map_location="cpu")
123
+ t2s_config = dict_s1["config"]
124
+ hz = 50
125
+ max_sec = t2s_config["data"]["max_sec"]
126
+
127
+ vq_model = SynthesizerTrn(
128
+ hps.data.filter_length // 2 + 1,
129
+ hps.train.segment_size // hps.data.hop_length,
130
+ n_speakers=hps.data.n_speakers,
131
+ **hps.model,
132
+ )
133
+ vq_model.load_state_dict(dict_s2["weight"], strict=False)
134
+
135
+ t2s_model = Text2SemanticLightningModule(t2s_config, "****", is_train=False)
136
+ t2s_model.load_state_dict(dict_s1["weight"])
137
+
138
+ if self.config.is_half:
139
+ bert_model = bert_model.half()
140
+ ssl_model = ssl_model.half()
141
+ vq_model = vq_model.half()
142
+ t2s_model = t2s_model.half()
143
+
144
+ t2s_model.model = torch.jit.script(t2s_model.model)
145
+
146
+ bert_model = bert_model.to(self.config.device)
147
+ ssl_model = ssl_model.to(self.config.device)
148
+ vq_model = vq_model.to(self.config.device)
149
+ t2s_model = t2s_model.to(self.config.device)
150
+ t2s_model.model = t2s_model.model.to(self.config.device)
151
+
152
+ bert_model.eval()
153
+ ssl_model.eval()
154
+ vq_model.eval()
155
+ t2s_model.eval()
156
+
157
+ self.tokenizer = tokenizer
158
+ self.bert_model = bert_model
159
+ self.ssl_model = ssl_model
160
+ self.vq_model = vq_model
161
+ self.t2s_model = t2s_model
162
+ self.hz = hz
163
+ self.max_sec = max_sec
164
+ self.hps = hps
165
+ self.zero_wav = np.zeros(
166
+ int(hps.data.sampling_rate * 0.3), dtype=self.config.numpy_dtype
167
+ )
168
+
169
+ def build_speaker(self, ref_audio, ref_text, ref_language) -> Speaker:
170
+ prompt_text = ref_text.strip("\n")
171
+ prompt_language = ref_language
172
+ with torch.no_grad():
173
+ wav16k = librosa.resample(ref_audio, orig_sr=32000, target_sr=16000)
174
+ wav16k = torch.from_numpy(wav16k)
175
+ zero_wav_torch = torch.from_numpy(self.zero_wav)
176
+ if self.config.is_half:
177
+ wav16k = wav16k.half().to(self.config.device)
178
+ zero_wav_torch = zero_wav_torch.half().to(self.config.device)
179
+ else:
180
+ wav16k = wav16k.to(self.config.device)
181
+ zero_wav_torch = zero_wav_torch.to(self.config.device)
182
+ wav16k_concat = torch.cat([wav16k, zero_wav_torch])
183
+ ssl_content = self.ssl_model.model(wav16k_concat.unsqueeze(0))[
184
+ "last_hidden_state"
185
+ ].transpose(
186
+ 1, 2
187
+ ) # .float()
188
+ codes = self.vq_model.extract_latent(ssl_content)
189
+ prompt_semantic = codes[0, 0]
190
+
191
+ refer = preprocess.get_spec(self.hps, ref_audio, 32000) # .to(device)
192
+ if self.config.is_half:
193
+ refer = refer.half()
194
+ refer = refer.to(self.config.device)
195
+
196
+ phones1, word2ph1, norm_text1 = preprocess.get_cleaned_text_final(
197
+ prompt_text, prompt_language
198
+ )
199
+ bert1 = self.get_bert_final(
200
+ phones1, word2ph1, norm_text1, prompt_language, self.config.device
201
+ ).to(self.config.torch_dtype)
202
+ return Speaker(refer, prompt_semantic, phones1, bert1)
203
+
204
+ def get_bert_feature(self, text, word2ph):
205
+ with torch.no_grad():
206
+ inputs = self.tokenizer(text, return_tensors="pt")
207
+ for i in inputs:
208
+ inputs[i] = inputs[i].to(
209
+ self.config.device
210
+ ) # 输入是long不用管精度问题,精度随bert_model
211
+ res = self.bert_model(**inputs, output_hidden_states=True)
212
+ res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1]
213
+ assert len(word2ph) == len(text)
214
+ phone_level_feature = []
215
+ for i in range(len(word2ph)):
216
+ repeat_feature = res[i].repeat(word2ph[i], 1)
217
+ phone_level_feature.append(repeat_feature)
218
+ phone_level_feature = torch.cat(phone_level_feature, dim=0)
219
+ # if(is_half==True):phone_level_feature=phone_level_feature.half()
220
+ return phone_level_feature.T
221
+
222
+ def get_bert_inf(self, phones, word2ph, norm_text, language):
223
+ language = language.replace("all_", "")
224
+ if language == "zh":
225
+ bert = self.get_bert_feature(norm_text, word2ph).to(
226
+ self.config.device
227
+ ) # .to(dtype)
228
+ else:
229
+ bert = torch.zeros((1024, len(phones)), dtype=self.config.torch_dtype).to(
230
+ self.config.device
231
+ )
232
+
233
+ return bert
234
+
235
+ def nonen_get_bert_inf(self, text, language):
236
+ if language != "auto":
237
+ textlist, langlist = preprocess.splite_en_inf(text, language)
238
+ else:
239
+ textlist = []
240
+ langlist = []
241
+ num_zh = 0
242
+ num_ja = 0
243
+ with preprocess.process_lock:
244
+ for tmp in LangSegment.getTexts(text):
245
+ langlist.append(tmp["lang"])
246
+ textlist.append(tmp["text"])
247
+ if tmp["lang"] == "zh":
248
+ num_zh += len(tmp["text"])
249
+ if tmp["lang"] == "ja":
250
+ num_ja += len(tmp["text"])
251
+
252
+ # Hotfix(yuheng): The LangSegment library has high error ratio.
253
+ # Apply quick fix to zh/ja.
254
+ if num_zh > 0 and num_ja > 0:
255
+ lang = "zh" if num_zh > 0 else "ja"
256
+ for i in range(len(langlist)):
257
+ if langlist[i] in ["zh", "ja"]:
258
+ langlist[i] = lang
259
+
260
+ bert_list = []
261
+ for i in range(len(textlist)):
262
+ lang = langlist[i]
263
+ phones, word2ph, norm_text = preprocess.clean_text_inf(textlist[i], lang)
264
+ bert = self.get_bert_inf(phones, word2ph, norm_text, lang)
265
+ bert_list.append(bert)
266
+ bert = torch.cat(bert_list, dim=1)
267
+
268
+ return bert
269
+
270
+ def get_bert_final(self, phones, word2ph, text, language, device):
271
+ if language == "en":
272
+ bert = self.get_bert_inf(phones, word2ph, text, language)
273
+ elif language in {"zh", "ja", "auto"}:
274
+ bert = self.nonen_get_bert_inf(text, language)
275
+ elif language == "all_zh":
276
+ bert = self.get_bert_feature(text, word2ph).to(device)
277
+ else:
278
+ bert = torch.zeros((1024, len(phones))).to(device)
279
+ return bert
280
+
281
+ def get_sample_rate(self):
282
+ return self.hps.data.sampling_rate
283
+
284
+ def has_meaningful_content(self, text):
285
+ # 去除字符串两端的空白字符
286
+ text = text.strip()
287
+
288
+ # 定义一个正则表达式匹配只包含标点符号、空格等的字符串
289
+ # 包括中文、日文、韩文和西文标点符号
290
+ pattern = re.compile(
291
+ f"^[{string.punctuation}、。!?,;:「」『』()[]【】〔〕ー-~。]+$"
292
+ )
293
+
294
+ # 检查字符串是否为空或只包含标点符号
295
+ if not text or pattern.match(text):
296
+ return False
297
+ else:
298
+ return True
299
+
300
+ def encode_refer(self, ref_audio, sr):
301
+ ref_spec = preprocess.get_spec_tensor(self.hps, ref_audio, sr).to("cuda").half()
302
+ return self.vq_model.ref_enc(ref_spec)
303
+
304
+ def _inference(
305
+ self,
306
+ speaker: Speaker,
307
+ text,
308
+ language,
309
+ top_k=5,
310
+ top_p=1,
311
+ temperature=1.0,
312
+ ge=None,
313
+ ):
314
+ if len(text.strip()) == 0:
315
+ return self.zero_wav
316
+
317
+ if text[-1] not in preprocess.splits:
318
+ text += "。" if language != "en" else "."
319
+
320
+ phones2, word2ph2, norm_text2 = preprocess.get_cleaned_text_final(
321
+ text, language
322
+ )
323
+ bert2 = self.get_bert_final(
324
+ phones2, word2ph2, norm_text2, language, self.config.device
325
+ ).to(self.config.torch_dtype)
326
+
327
+ bert = torch.cat([speaker.bert, bert2], 1)
328
+ all_phoneme_ids = (
329
+ torch.LongTensor(speaker.phoneme + phones2)
330
+ .to(self.config.device)
331
+ .unsqueeze(0)
332
+ )
333
+ bert = bert.to(self.config.device).unsqueeze(0)
334
+ prompt = speaker.prompt_semantic.unsqueeze(0).to(self.config.device)
335
+
336
+ with torch.no_grad():
337
+ pred_semantic, idx = self.t2s_model.model(
338
+ all_phoneme_ids.clone(),
339
+ prompt.clone(),
340
+ bert.clone(),
341
+ top_k=top_k,
342
+ top_p=top_p,
343
+ temperature=temperature,
344
+ early_stop_num=self.hz * self.max_sec,
345
+ )
346
+
347
+ pred_semantic = pred_semantic[:, -idx:].unsqueeze(
348
+ 0
349
+ ) # .unsqueeze(0)#mq要多unsqueeze一次
350
+
351
+ if ge is not None:
352
+ audio = self.vq_model.decode_by_ge(
353
+ pred_semantic,
354
+ torch.LongTensor(phones2).to(self.config.device).unsqueeze(0),
355
+ ge,
356
+ )
357
+ else:
358
+ audio = self.vq_model.decode(
359
+ pred_semantic,
360
+ torch.LongTensor(phones2).to(self.config.device).unsqueeze(0),
361
+ speaker.refer,
362
+ )
363
+
364
+ audio = audio.detach().cpu().numpy()[0, 0]
365
+
366
+ max_audio = np.abs(audio).max() # 简单防止16bit爆音
367
+ if max_audio > 1:
368
+ audio /= max_audio
369
+
370
+ return audio
371
+
372
+ def inference(
373
+ self,
374
+ speaker: Speaker,
375
+ text,
376
+ language,
377
+ top_k=5,
378
+ top_p=1,
379
+ temperature=1.0,
380
+ ge=None,
381
+ ):
382
+ while "\n\n" in text:
383
+ text = text.replace("\n\n", "\n")
384
+
385
+ text = emoji.replace_emoji(text, "")
386
+ texts = text.split("\n")
387
+ texts = [text.strip() for text in texts if len(text.strip()) > 0]
388
+ texts = preprocess.merge_short_text_in_array(texts, 5)
389
+
390
+ if language == "auto":
391
+ language = langid.classify(text)[0]
392
+
393
+ audio_opt = []
394
+
395
+ for text in texts:
396
+ if len(text.strip()) == 0:
397
+ continue
398
+
399
+ if not self.has_meaningful_content(text):
400
+ continue
401
+
402
+ # if the whole text is composed of numbers, we will skip it, because we can't detect the language of the text
403
+ if language == "auto" and re.match(r"^[\d\.,:/#]+$", text):
404
+ continue
405
+
406
+ audio = self._inference(
407
+ speaker,
408
+ text,
409
+ language,
410
+ top_k,
411
+ top_p,
412
+ temperature,
413
+ ge=ge,
414
+ )
415
+ audio_opt.append(audio)
416
+ audio_opt.append(self.zero_wav)
417
+
418
+ if len(audio_opt) == 0:
419
+ return self.zero_wav
420
+ audio = np.concatenate(audio_opt)
421
+ return audio
422
+
423
+ def inference_to_file(
424
+ self,
425
+ speaker: Speaker,
426
+ text,
427
+ language,
428
+ output_path,
429
+ top_k=5,
430
+ top_p=1,
431
+ temperature=1.0,
432
+ ge=None,
433
+ ):
434
+ audio = self.inference(
435
+ speaker,
436
+ text,
437
+ language,
438
+ top_k,
439
+ top_p,
440
+ temperature,
441
+ ge=ge,
442
+ )
443
+ if self.config.is_half:
444
+ audio = audio.astype(np.float32)
445
+ sf.write(output_path, audio, self.get_sample_rate())
lipsync-models/GPT_SoVITS/inference_gui.py ADDED
@@ -0,0 +1,340 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+ from PyQt5.QtCore import QEvent
3
+ from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QLineEdit, QPushButton, QTextEdit
4
+ from PyQt5.QtWidgets import QGridLayout, QVBoxLayout, QWidget, QFileDialog, QStatusBar, QComboBox
5
+ import soundfile as sf
6
+
7
+ from tools.i18n.i18n import I18nAuto
8
+ i18n = I18nAuto()
9
+
10
+ from GPT_SoVITS.inference_webui import change_gpt_weights, change_sovits_weights, get_tts_wav
11
+
12
+
13
+ class GPTSoVITSGUI(QMainWindow):
14
+ def __init__(self):
15
+ super().__init__()
16
+
17
+ self.init_ui()
18
+
19
+ def init_ui(self):
20
+ self.setWindowTitle('GPT-SoVITS GUI')
21
+ self.setGeometry(800, 450, 950, 850)
22
+
23
+ self.setStyleSheet("""
24
+ QWidget {
25
+ background-color: #a3d3b1;
26
+ }
27
+
28
+ QTabWidget::pane {
29
+ background-color: #a3d3b1;
30
+ }
31
+
32
+ QTabWidget::tab-bar {
33
+ alignment: left;
34
+ }
35
+
36
+ QTabBar::tab {
37
+ background: #8da4bf;
38
+ color: #ffffff;
39
+ padding: 8px;
40
+ }
41
+
42
+ QTabBar::tab:selected {
43
+ background: #2a3f54;
44
+ }
45
+
46
+ QLabel {
47
+ color: #000000;
48
+ }
49
+
50
+ QPushButton {
51
+ background-color: #4CAF50;
52
+ color: white;
53
+ padding: 8px;
54
+ border: 1px solid #4CAF50;
55
+ border-radius: 4px;
56
+ }
57
+
58
+ QPushButton:hover {
59
+ background-color: #45a049;
60
+ border: 1px solid #45a049;
61
+ box-shadow: 2px 2px 2px rgba(0, 0, 0, 0.1);
62
+ }
63
+ """)
64
+
65
+ license_text = (
66
+ "本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. "
67
+ "如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录LICENSE.")
68
+ license_label = QLabel(license_text)
69
+ license_label.setWordWrap(True)
70
+
71
+ self.GPT_model_label = QLabel("选择GPT模型:")
72
+ self.GPT_model_input = QLineEdit()
73
+ self.GPT_model_input.setPlaceholderText("拖拽或选择文件")
74
+ self.GPT_model_input.setReadOnly(True)
75
+ self.GPT_model_button = QPushButton("选择GPT模型文件")
76
+ self.GPT_model_button.clicked.connect(self.select_GPT_model)
77
+
78
+ self.SoVITS_model_label = QLabel("选择SoVITS模型:")
79
+ self.SoVITS_model_input = QLineEdit()
80
+ self.SoVITS_model_input.setPlaceholderText("拖拽或选择文件")
81
+ self.SoVITS_model_input.setReadOnly(True)
82
+ self.SoVITS_model_button = QPushButton("选择SoVITS模型文件")
83
+ self.SoVITS_model_button.clicked.connect(self.select_SoVITS_model)
84
+
85
+ self.ref_audio_label = QLabel("上传参考音频:")
86
+ self.ref_audio_input = QLineEdit()
87
+ self.ref_audio_input.setPlaceholderText("拖拽或选择文件")
88
+ self.ref_audio_input.setReadOnly(True)
89
+ self.ref_audio_button = QPushButton("选择音频文件")
90
+ self.ref_audio_button.clicked.connect(self.select_ref_audio)
91
+
92
+ self.ref_text_label = QLabel("参考音频文本:")
93
+ self.ref_text_input = QLineEdit()
94
+ self.ref_text_input.setPlaceholderText("拖拽或选择文件")
95
+ self.ref_text_input.setReadOnly(True)
96
+ self.ref_text_button = QPushButton("上传文本")
97
+ self.ref_text_button.clicked.connect(self.upload_ref_text)
98
+
99
+ self.language_label = QLabel("参考音频语言:")
100
+ self.language_combobox = QComboBox()
101
+ self.language_combobox.addItems(["中文", "英文", "日文"])
102
+
103
+ self.target_text_label = QLabel("合成目标文本:")
104
+ self.target_text_input = QLineEdit()
105
+ self.target_text_input.setPlaceholderText("拖拽或选择文件")
106
+ self.target_text_input.setReadOnly(True)
107
+ self.target_text_button = QPushButton("上传文本")
108
+ self.target_text_button.clicked.connect(self.upload_target_text)
109
+
110
+ self.language_label_02 = QLabel("合成音频语言:")
111
+ self.language_combobox_02 = QComboBox()
112
+ self.language_combobox_02.addItems(["中文", "英文", "日文"])
113
+
114
+ self.output_label = QLabel("输出音频路径:")
115
+ self.output_input = QLineEdit()
116
+ self.output_input.setPlaceholderText("拖拽或选择文件")
117
+ self.output_input.setReadOnly(True)
118
+ self.output_button = QPushButton("选择文件夹")
119
+ self.output_button.clicked.connect(self.select_output_path)
120
+
121
+ self.output_text = QTextEdit()
122
+ self.output_text.setReadOnly(True)
123
+
124
+ self.add_drag_drop_events([
125
+ self.GPT_model_input,
126
+ self.SoVITS_model_input,
127
+ self.ref_audio_input,
128
+ self.ref_text_input,
129
+ self.target_text_input,
130
+ self.output_input,
131
+ ])
132
+
133
+ self.synthesize_button = QPushButton("合成")
134
+ self.synthesize_button.clicked.connect(self.synthesize)
135
+
136
+ self.clear_output_button = QPushButton("清空输出")
137
+ self.clear_output_button.clicked.connect(self.clear_output)
138
+
139
+ self.status_bar = QStatusBar()
140
+
141
+ main_layout = QVBoxLayout()
142
+
143
+ input_layout = QGridLayout()
144
+ input_layout.setSpacing(10)
145
+
146
+ self.setLayout(input_layout)
147
+
148
+ input_layout.addWidget(license_label, 0, 0, 1, 3)
149
+
150
+ input_layout.addWidget(self.GPT_model_label, 1, 0)
151
+ input_layout.addWidget(self.GPT_model_input, 2, 0, 1, 2)
152
+ input_layout.addWidget(self.GPT_model_button, 2, 2)
153
+
154
+ input_layout.addWidget(self.SoVITS_model_label, 3, 0)
155
+ input_layout.addWidget(self.SoVITS_model_input, 4, 0, 1, 2)
156
+ input_layout.addWidget(self.SoVITS_model_button, 4, 2)
157
+
158
+ input_layout.addWidget(self.ref_audio_label, 5, 0)
159
+ input_layout.addWidget(self.ref_audio_input, 6, 0, 1, 2)
160
+ input_layout.addWidget(self.ref_audio_button, 6, 2)
161
+
162
+ input_layout.addWidget(self.language_label, 7, 0)
163
+ input_layout.addWidget(self.language_combobox, 8, 0, 1, 1)
164
+ input_layout.addWidget(self.ref_text_label, 9, 0)
165
+ input_layout.addWidget(self.ref_text_input, 10, 0, 1, 2)
166
+ input_layout.addWidget(self.ref_text_button, 10, 2)
167
+
168
+ input_layout.addWidget(self.language_label_02, 11, 0)
169
+ input_layout.addWidget(self.language_combobox_02, 12, 0, 1, 1)
170
+ input_layout.addWidget(self.target_text_label, 13, 0)
171
+ input_layout.addWidget(self.target_text_input, 14, 0, 1, 2)
172
+ input_layout.addWidget(self.target_text_button, 14, 2)
173
+
174
+ input_layout.addWidget(self.output_label, 15, 0)
175
+ input_layout.addWidget(self.output_input, 16, 0, 1, 2)
176
+ input_layout.addWidget(self.output_button, 16, 2)
177
+
178
+ main_layout.addLayout(input_layout)
179
+
180
+ output_layout = QVBoxLayout()
181
+ output_layout.addWidget(self.output_text)
182
+ main_layout.addLayout(output_layout)
183
+
184
+ main_layout.addWidget(self.synthesize_button)
185
+
186
+ main_layout.addWidget(self.clear_output_button)
187
+
188
+ main_layout.addWidget(self.status_bar)
189
+
190
+ self.central_widget = QWidget()
191
+ self.central_widget.setLayout(main_layout)
192
+ self.setCentralWidget(self.central_widget)
193
+
194
+ def dragEnterEvent(self, event):
195
+ if event.mimeData().hasUrls():
196
+ event.acceptProposedAction()
197
+
198
+ def dropEvent(self, event):
199
+ if event.mimeData().hasUrls():
200
+ file_paths = [url.toLocalFile() for url in event.mimeData().urls()]
201
+
202
+ if len(file_paths) == 1:
203
+ self.update_ref_audio(file_paths[0])
204
+ self.update_input_paths(self.ref_audio_input, file_paths[0])
205
+ else:
206
+ self.update_ref_audio(", ".join(file_paths))
207
+
208
+ def add_drag_drop_events(self, widgets):
209
+ for widget in widgets:
210
+ widget.setAcceptDrops(True)
211
+ widget.installEventFilter(self)
212
+
213
+ def eventFilter(self, obj, event):
214
+ if event.type() == QEvent.DragEnter:
215
+ mime_data = event.mimeData()
216
+ if mime_data.hasUrls():
217
+ event.acceptProposedAction()
218
+
219
+ elif event.type() == QEvent.Drop:
220
+ mime_data = event.mimeData()
221
+ if mime_data.hasUrls():
222
+ file_paths = [url.toLocalFile() for url in mime_data.urls()]
223
+ if len(file_paths) == 1:
224
+ self.update_input_paths(obj, file_paths[0])
225
+ else:
226
+ self.update_input_paths(obj, ", ".join(file_paths))
227
+ event.acceptProposedAction()
228
+
229
+ return super().eventFilter(obj, event)
230
+
231
+ def select_GPT_model(self):
232
+ file_path, _ = QFileDialog.getOpenFileName(self, "选择GPT模型文件", "", "GPT Files (*.ckpt)")
233
+ if file_path:
234
+ self.GPT_model_input.setText(file_path)
235
+
236
+ def select_SoVITS_model(self):
237
+ file_path, _ = QFileDialog.getOpenFileName(self, "选择SoVITS模型文件", "", "SoVITS Files (*.pth)")
238
+ if file_path:
239
+ self.SoVITS_model_input.setText(file_path)
240
+
241
+ def select_ref_audio(self):
242
+ options = QFileDialog.Options()
243
+ options |= QFileDialog.DontUseNativeDialog
244
+ options |= QFileDialog.ShowDirsOnly
245
+
246
+ file_dialog = QFileDialog()
247
+ file_dialog.setOptions(options)
248
+
249
+ file_dialog.setFileMode(QFileDialog.AnyFile)
250
+ file_dialog.setNameFilter("Audio Files (*.wav *.mp3)")
251
+
252
+ if file_dialog.exec_():
253
+ file_paths = file_dialog.selectedFiles()
254
+
255
+ if len(file_paths) == 1:
256
+ self.update_ref_audio(file_paths[0])
257
+ self.update_input_paths(self.ref_audio_input, file_paths[0])
258
+ else:
259
+ self.update_ref_audio(", ".join(file_paths))
260
+
261
+ def upload_ref_text(self):
262
+ file_path, _ = QFileDialog.getOpenFileName(self, "选择文本文件", "", "Text Files (*.txt)")
263
+ if file_path:
264
+ with open(file_path, 'r', encoding='utf-8') as file:
265
+ content = file.read()
266
+ self.ref_text_input.setText(content)
267
+ self.update_input_paths(self.ref_text_input, file_path)
268
+
269
+ def upload_target_text(self):
270
+ file_path, _ = QFileDialog.getOpenFileName(self, "选择文本文件", "", "Text Files (*.txt)")
271
+ if file_path:
272
+ with open(file_path, 'r', encoding='utf-8') as file:
273
+ content = file.read()
274
+ self.target_text_input.setText(content)
275
+ self.update_input_paths(self.target_text_input, file_path)
276
+
277
+ def select_output_path(self):
278
+ options = QFileDialog.Options()
279
+ options |= QFileDialog.DontUseNativeDialog
280
+ options |= QFileDialog.ShowDirsOnly
281
+
282
+ folder_dialog = QFileDialog()
283
+ folder_dialog.setOptions(options)
284
+ folder_dialog.setFileMode(QFileDialog.Directory)
285
+
286
+ if folder_dialog.exec_():
287
+ folder_path = folder_dialog.selectedFiles()[0]
288
+ self.output_input.setText(folder_path)
289
+
290
+ def update_ref_audio(self, file_path):
291
+ self.ref_audio_input.setText(file_path)
292
+
293
+ def update_input_paths(self, input_box, file_path):
294
+ input_box.setText(file_path)
295
+
296
+ def clear_output(self):
297
+ self.output_text.clear()
298
+
299
+ def synthesize(self):
300
+ GPT_model_path = self.GPT_model_input.text()
301
+ SoVITS_model_path = self.SoVITS_model_input.text()
302
+ ref_audio_path = self.ref_audio_input.text()
303
+ language_combobox = self.language_combobox.currentText()
304
+ language_combobox = i18n(language_combobox)
305
+ ref_text = self.ref_text_input.text()
306
+ language_combobox_02 = self.language_combobox_02.currentText()
307
+ language_combobox_02 = i18n(language_combobox_02)
308
+ target_text = self.target_text_input.text()
309
+ output_path = self.output_input.text()
310
+
311
+ change_gpt_weights(gpt_path=GPT_model_path)
312
+ change_sovits_weights(sovits_path=SoVITS_model_path)
313
+
314
+ synthesis_result = get_tts_wav(ref_wav_path=ref_audio_path,
315
+ prompt_text=ref_text,
316
+ prompt_language=language_combobox,
317
+ text=target_text,
318
+ text_language=language_combobox_02)
319
+
320
+ result_list = list(synthesis_result)
321
+
322
+ if result_list:
323
+ last_sampling_rate, last_audio_data = result_list[-1]
324
+ output_wav_path = os.path.join(output_path, "output.wav")
325
+ sf.write(output_wav_path, last_audio_data, last_sampling_rate)
326
+
327
+ result = "Audio saved to " + output_wav_path
328
+
329
+ self.status_bar.showMessage("合成完成!输出路径:" + output_wav_path, 5000)
330
+ self.output_text.append("处理结果:\n" + result)
331
+
332
+ def main():
333
+ app = QApplication(sys.argv)
334
+ mainWin = GPTSoVITSGUI()
335
+ mainWin.show()
336
+ sys.exit(app.exec_())
337
+
338
+
339
+ if __name__ == '__main__':
340
+ main()
lipsync-models/GPT_SoVITS/inference_webui.py ADDED
@@ -0,0 +1,679 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ '''
2
+ 按中英混合识别
3
+ 按日英混合识别
4
+ 多语种启动切分识别语种
5
+ 全部按中文识别
6
+ 全部按英文识别
7
+ 全部按日文识别
8
+ '''
9
+ import os, re, logging
10
+ import LangSegment
11
+ logging.getLogger("markdown_it").setLevel(logging.ERROR)
12
+ logging.getLogger("urllib3").setLevel(logging.ERROR)
13
+ logging.getLogger("httpcore").setLevel(logging.ERROR)
14
+ logging.getLogger("httpx").setLevel(logging.ERROR)
15
+ logging.getLogger("asyncio").setLevel(logging.ERROR)
16
+ logging.getLogger("charset_normalizer").setLevel(logging.ERROR)
17
+ logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
18
+ import pdb
19
+
20
+ if os.path.exists("./gweight.txt"):
21
+ with open("./gweight.txt", 'r', encoding="utf-8") as file:
22
+ gweight_data = file.read()
23
+ gpt_path = os.environ.get(
24
+ "gpt_path", gweight_data)
25
+ else:
26
+ gpt_path = os.environ.get(
27
+ "gpt_path", "GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt")
28
+
29
+ if os.path.exists("./sweight.txt"):
30
+ with open("./sweight.txt", 'r', encoding="utf-8") as file:
31
+ sweight_data = file.read()
32
+ sovits_path = os.environ.get("sovits_path", sweight_data)
33
+ else:
34
+ sovits_path = os.environ.get("sovits_path", "GPT_SoVITS/pretrained_models/s2G488k.pth")
35
+ # gpt_path = os.environ.get(
36
+ # "gpt_path", "pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"
37
+ # )
38
+ # sovits_path = os.environ.get("sovits_path", "pretrained_models/s2G488k.pth")
39
+ cnhubert_base_path = os.environ.get(
40
+ "cnhubert_base_path", "GPT_SoVITS/pretrained_models/chinese-hubert-base"
41
+ )
42
+ bert_path = os.environ.get(
43
+ "bert_path", "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large"
44
+ )
45
+ infer_ttswebui = os.environ.get("infer_ttswebui", 9872)
46
+ infer_ttswebui = int(infer_ttswebui)
47
+ is_share = os.environ.get("is_share", "False")
48
+ is_share = eval(is_share)
49
+ if "_CUDA_VISIBLE_DEVICES" in os.environ:
50
+ os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["_CUDA_VISIBLE_DEVICES"]
51
+ is_half = eval(os.environ.get("is_half", "True"))
52
+ import gradio as gr
53
+ from transformers import AutoModelForMaskedLM, AutoTokenizer
54
+ import numpy as np
55
+ import librosa, torch
56
+ from feature_extractor import cnhubert
57
+
58
+ cnhubert.cnhubert_base_path = cnhubert_base_path
59
+
60
+ from module.models import SynthesizerTrn
61
+ from AR.models.t2s_lightning_module import Text2SemanticLightningModule
62
+ from text import cleaned_text_to_sequence
63
+ from text.cleaner import clean_text
64
+ from time import time as ttime
65
+ from module.mel_processing import spectrogram_torch
66
+ from my_utils import load_audio
67
+ from tools.i18n.i18n import I18nAuto
68
+
69
+ i18n = I18nAuto()
70
+
71
+ os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # 确保直接启动推理UI时也能够设置。
72
+
73
+ if torch.cuda.is_available():
74
+ device = "cuda"
75
+ elif torch.backends.mps.is_available():
76
+ device = "mps"
77
+ else:
78
+ device = "cpu"
79
+
80
+ tokenizer = AutoTokenizer.from_pretrained(bert_path)
81
+ bert_model = AutoModelForMaskedLM.from_pretrained(bert_path)
82
+ if is_half == True:
83
+ bert_model = bert_model.half().to(device)
84
+ else:
85
+ bert_model = bert_model.to(device)
86
+
87
+
88
+ def get_bert_feature(text, word2ph):
89
+ with torch.no_grad():
90
+ inputs = tokenizer(text, return_tensors="pt")
91
+ for i in inputs:
92
+ inputs[i] = inputs[i].to(device)
93
+ res = bert_model(**inputs, output_hidden_states=True)
94
+ res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1]
95
+ assert len(word2ph) == len(text)
96
+ phone_level_feature = []
97
+ for i in range(len(word2ph)):
98
+ repeat_feature = res[i].repeat(word2ph[i], 1)
99
+ phone_level_feature.append(repeat_feature)
100
+ phone_level_feature = torch.cat(phone_level_feature, dim=0)
101
+ return phone_level_feature.T
102
+
103
+
104
+ class DictToAttrRecursive(dict):
105
+ def __init__(self, input_dict):
106
+ super().__init__(input_dict)
107
+ for key, value in input_dict.items():
108
+ if isinstance(value, dict):
109
+ value = DictToAttrRecursive(value)
110
+ self[key] = value
111
+ setattr(self, key, value)
112
+
113
+ def __getattr__(self, item):
114
+ try:
115
+ return self[item]
116
+ except KeyError:
117
+ raise AttributeError(f"Attribute {item} not found")
118
+
119
+ def __setattr__(self, key, value):
120
+ if isinstance(value, dict):
121
+ value = DictToAttrRecursive(value)
122
+ super(DictToAttrRecursive, self).__setitem__(key, value)
123
+ super().__setattr__(key, value)
124
+
125
+ def __delattr__(self, item):
126
+ try:
127
+ del self[item]
128
+ except KeyError:
129
+ raise AttributeError(f"Attribute {item} not found")
130
+
131
+
132
+ ssl_model = cnhubert.get_model()
133
+ if is_half == True:
134
+ ssl_model = ssl_model.half().to(device)
135
+ else:
136
+ ssl_model = ssl_model.to(device)
137
+
138
+
139
+ def change_sovits_weights(sovits_path):
140
+ global vq_model, hps
141
+ dict_s2 = torch.load(sovits_path, map_location="cpu")
142
+ hps = dict_s2["config"]
143
+ hps = DictToAttrRecursive(hps)
144
+ hps.model.semantic_frame_rate = "25hz"
145
+ vq_model = SynthesizerTrn(
146
+ hps.data.filter_length // 2 + 1,
147
+ hps.train.segment_size // hps.data.hop_length,
148
+ n_speakers=hps.data.n_speakers,
149
+ **hps.model
150
+ )
151
+ if ("pretrained" not in sovits_path):
152
+ del vq_model.enc_q
153
+ if is_half == True:
154
+ vq_model = vq_model.half().to(device)
155
+ else:
156
+ vq_model = vq_model.to(device)
157
+ vq_model.eval()
158
+ print(vq_model.load_state_dict(dict_s2["weight"], strict=False))
159
+ with open("./sweight.txt", "w", encoding="utf-8") as f:
160
+ f.write(sovits_path)
161
+
162
+
163
+ change_sovits_weights(sovits_path)
164
+
165
+
166
+ def change_gpt_weights(gpt_path):
167
+ global hz, max_sec, t2s_model, config
168
+ hz = 50
169
+ dict_s1 = torch.load(gpt_path, map_location="cpu")
170
+ config = dict_s1["config"]
171
+ max_sec = config["data"]["max_sec"]
172
+ t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
173
+ t2s_model.load_state_dict(dict_s1["weight"])
174
+ if is_half == True:
175
+ t2s_model = t2s_model.half()
176
+ t2s_model = t2s_model.to(device)
177
+ t2s_model.eval()
178
+ total = sum([param.nelement() for param in t2s_model.parameters()])
179
+ print("Number of parameter: %.2fM" % (total / 1e6))
180
+ with open("./gweight.txt", "w", encoding="utf-8") as f: f.write(gpt_path)
181
+
182
+
183
+ change_gpt_weights(gpt_path)
184
+
185
+
186
+ def get_spepc(hps, filename):
187
+ audio = load_audio(filename, int(hps.data.sampling_rate))
188
+ audio = torch.FloatTensor(audio)
189
+ audio_norm = audio
190
+ audio_norm = audio_norm.unsqueeze(0)
191
+ spec = spectrogram_torch(
192
+ audio_norm,
193
+ hps.data.filter_length,
194
+ hps.data.sampling_rate,
195
+ hps.data.hop_length,
196
+ hps.data.win_length,
197
+ center=False,
198
+ )
199
+ return spec
200
+
201
+
202
+ dict_language = {
203
+ i18n("中文"): "all_zh",#全部按中文识别
204
+ i18n("英文"): "en",#全部按英文识别#######不变
205
+ i18n("日文"): "all_ja",#全部按日文识别
206
+ i18n("中英混合"): "zh",#按中英混合识别####不变
207
+ i18n("日英混合"): "ja",#按日英混合识别####不变
208
+ i18n("多语种混合"): "auto",#多语种启动切分识别语种
209
+ }
210
+
211
+
212
+ def splite_en_inf(sentence, language):
213
+ pattern = re.compile(r'[a-zA-Z ]+')
214
+ textlist = []
215
+ langlist = []
216
+ pos = 0
217
+ for match in pattern.finditer(sentence):
218
+ start, end = match.span()
219
+ if start > pos:
220
+ textlist.append(sentence[pos:start])
221
+ langlist.append(language)
222
+ textlist.append(sentence[start:end])
223
+ langlist.append("en")
224
+ pos = end
225
+ if pos < len(sentence):
226
+ textlist.append(sentence[pos:])
227
+ langlist.append(language)
228
+ # Merge punctuation into previous word
229
+ for i in range(len(textlist)-1, 0, -1):
230
+ if re.match(r'^[\W_]+$', textlist[i]):
231
+ textlist[i-1] += textlist[i]
232
+ del textlist[i]
233
+ del langlist[i]
234
+ # Merge consecutive words with the same language tag
235
+ i = 0
236
+ while i < len(langlist) - 1:
237
+ if langlist[i] == langlist[i+1]:
238
+ textlist[i] += textlist[i+1]
239
+ del textlist[i+1]
240
+ del langlist[i+1]
241
+ else:
242
+ i += 1
243
+
244
+ return textlist, langlist
245
+
246
+
247
+ def clean_text_inf(text, language):
248
+ formattext = ""
249
+ language = language.replace("all_","")
250
+ for tmp in LangSegment.getTexts(text):
251
+ if language == "ja":
252
+ if tmp["lang"] == language or tmp["lang"] == "zh":
253
+ formattext += tmp["text"] + " "
254
+ continue
255
+ if tmp["lang"] == language:
256
+ formattext += tmp["text"] + " "
257
+ while " " in formattext:
258
+ formattext = formattext.replace(" ", " ")
259
+ phones, word2ph, norm_text = clean_text(formattext, language)
260
+ phones = cleaned_text_to_sequence(phones)
261
+ return phones, word2ph, norm_text
262
+
263
+ dtype=torch.float16 if is_half == True else torch.float32
264
+ def get_bert_inf(phones, word2ph, norm_text, language):
265
+ language=language.replace("all_","")
266
+ if language == "zh":
267
+ bert = get_bert_feature(norm_text, word2ph).to(device)#.to(dtype)
268
+ else:
269
+ bert = torch.zeros(
270
+ (1024, len(phones)),
271
+ dtype=torch.float16 if is_half == True else torch.float32,
272
+ ).to(device)
273
+
274
+ return bert
275
+
276
+
277
+ def nonen_clean_text_inf(text, language):
278
+ if(language!="auto"):
279
+ textlist, langlist = splite_en_inf(text, language)
280
+ else:
281
+ textlist=[]
282
+ langlist=[]
283
+ for tmp in LangSegment.getTexts(text):
284
+ langlist.append(tmp["lang"])
285
+ textlist.append(tmp["text"])
286
+ phones_list = []
287
+ word2ph_list = []
288
+ norm_text_list = []
289
+ for i in range(len(textlist)):
290
+ lang = langlist[i]
291
+ phones, word2ph, norm_text = clean_text_inf(textlist[i], lang)
292
+ phones_list.append(phones)
293
+ if lang == "zh":
294
+ word2ph_list.append(word2ph)
295
+ norm_text_list.append(norm_text)
296
+ print(word2ph_list)
297
+ phones = sum(phones_list, [])
298
+ word2ph = sum(word2ph_list, [])
299
+ norm_text = ' '.join(norm_text_list)
300
+
301
+ return phones, word2ph, norm_text
302
+
303
+
304
+ def nonen_get_bert_inf(text, language):
305
+ if(language!="auto"):
306
+ textlist, langlist = splite_en_inf(text, language)
307
+ else:
308
+ textlist=[]
309
+ langlist=[]
310
+ for tmp in LangSegment.getTexts(text):
311
+ langlist.append(tmp["lang"])
312
+ textlist.append(tmp["text"])
313
+ print(textlist)
314
+ print(langlist)
315
+ bert_list = []
316
+ for i in range(len(textlist)):
317
+ lang = langlist[i]
318
+ phones, word2ph, norm_text = clean_text_inf(textlist[i], lang)
319
+ bert = get_bert_inf(phones, word2ph, norm_text, lang)
320
+ bert_list.append(bert)
321
+ bert = torch.cat(bert_list, dim=1)
322
+
323
+ return bert
324
+
325
+
326
+ splits = {",", "。", "?", "!", ",", ".", "?", "!", "~", ":", ":", "—", "…", }
327
+
328
+
329
+ def get_first(text):
330
+ pattern = "[" + "".join(re.escape(sep) for sep in splits) + "]"
331
+ text = re.split(pattern, text)[0].strip()
332
+ return text
333
+
334
+
335
+ def get_cleaned_text_final(text,language):
336
+ if language in {"en","all_zh","all_ja"}:
337
+ phones, word2ph, norm_text = clean_text_inf(text, language)
338
+ elif language in {"zh", "ja","auto"}:
339
+ phones, word2ph, norm_text = nonen_clean_text_inf(text, language)
340
+ return phones, word2ph, norm_text
341
+
342
+ def get_bert_final(phones, word2ph, text,language,device):
343
+ if language == "en":
344
+ bert = get_bert_inf(phones, word2ph, text, language)
345
+ elif language in {"zh", "ja","auto"}:
346
+ bert = nonen_get_bert_inf(text, language)
347
+ elif language == "all_zh":
348
+ bert = get_bert_feature(text, word2ph).to(device)
349
+ else:
350
+ bert = torch.zeros((1024, len(phones))).to(device)
351
+ return bert
352
+
353
+ def merge_short_text_in_array(texts, threshold):
354
+ if (len(texts)) < 2:
355
+ return texts
356
+ result = []
357
+ text = ""
358
+ for ele in texts:
359
+ text += ele
360
+ if len(text) >= threshold:
361
+ result.append(text)
362
+ text = ""
363
+ if (len(text) > 0):
364
+ if len(result) == 0:
365
+ result.append(text)
366
+ else:
367
+ result[len(result) - 1] += text
368
+ return result
369
+
370
+ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language, how_to_cut=i18n("不切"), top_k=20, top_p=0.6, temperature=0.6, ref_free = False):
371
+ if prompt_text is None or len(prompt_text) == 0:
372
+ ref_free = True
373
+ t0 = ttime()
374
+ prompt_language = dict_language[prompt_language]
375
+ text_language = dict_language[text_language]
376
+ if not ref_free:
377
+ prompt_text = prompt_text.strip("\n")
378
+ if (prompt_text[-1] not in splits): prompt_text += "。" if prompt_language != "en" else "."
379
+ print(i18n("实际输入的参考文本:"), prompt_text)
380
+ text = text.strip("\n")
381
+ if (text[0] not in splits and len(get_first(text)) < 4): text = "。" + text if text_language != "en" else "." + text
382
+
383
+ print(i18n("实际输入的目标文本:"), text)
384
+ zero_wav = np.zeros(
385
+ int(hps.data.sampling_rate * 0.3),
386
+ dtype=np.float16 if is_half == True else np.float32,
387
+ )
388
+ with torch.no_grad():
389
+ wav16k, sr = librosa.load(ref_wav_path, sr=16000)
390
+ if (wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000):
391
+ raise OSError(i18n("参考音频在3~10秒范围外,请更换!"))
392
+ wav16k = torch.from_numpy(wav16k)
393
+ zero_wav_torch = torch.from_numpy(zero_wav)
394
+ if is_half == True:
395
+ wav16k = wav16k.half().to(device)
396
+ zero_wav_torch = zero_wav_torch.half().to(device)
397
+ else:
398
+ wav16k = wav16k.to(device)
399
+ zero_wav_torch = zero_wav_torch.to(device)
400
+ wav16k = torch.cat([wav16k, zero_wav_torch])
401
+ ssl_content = ssl_model.model(wav16k.unsqueeze(0))[
402
+ "last_hidden_state"
403
+ ].transpose(
404
+ 1, 2
405
+ ) # .float()
406
+ codes = vq_model.extract_latent(ssl_content)
407
+
408
+ prompt_semantic = codes[0, 0]
409
+ t1 = ttime()
410
+
411
+ if (how_to_cut == i18n("凑四句一切")):
412
+ text = cut1(text)
413
+ elif (how_to_cut == i18n("凑50字一切")):
414
+ text = cut2(text)
415
+ elif (how_to_cut == i18n("按中文句号。切")):
416
+ text = cut3(text)
417
+ elif (how_to_cut == i18n("按英文句号.切")):
418
+ text = cut4(text)
419
+ elif (how_to_cut == i18n("按标点符号切")):
420
+ text = cut5(text)
421
+ while "\n\n" in text:
422
+ text = text.replace("\n\n", "\n")
423
+ print(i18n("实际输入的目标文本(切句后):"), text)
424
+ texts = text.split("\n")
425
+ texts = merge_short_text_in_array(texts, 5)
426
+ audio_opt = []
427
+ if not ref_free:
428
+ phones1, word2ph1, norm_text1=get_cleaned_text_final(prompt_text, prompt_language)
429
+ bert1=get_bert_final(phones1, word2ph1, norm_text1,prompt_language,device).to(dtype)
430
+
431
+ for text in texts:
432
+ # 解决输入目标文本的空行导致报错的问题
433
+ if (len(text.strip()) == 0):
434
+ continue
435
+ if (text[-1] not in splits): text += "。" if text_language != "en" else "."
436
+ print(i18n("实际输入的目标文本(每句):"), text)
437
+ phones2, word2ph2, norm_text2 = get_cleaned_text_final(text, text_language)
438
+ bert2 = get_bert_final(phones2, word2ph2, norm_text2, text_language, device).to(dtype)
439
+ if not ref_free:
440
+ bert = torch.cat([bert1, bert2], 1)
441
+ all_phoneme_ids = torch.LongTensor(phones1+phones2).to(device).unsqueeze(0)
442
+ else:
443
+ bert = bert2
444
+ all_phoneme_ids = torch.LongTensor(phones2).to(device).unsqueeze(0)
445
+
446
+ bert = bert.to(device).unsqueeze(0)
447
+ all_phoneme_len = torch.tensor([all_phoneme_ids.shape[-1]]).to(device)
448
+ prompt = prompt_semantic.unsqueeze(0).to(device)
449
+ t2 = ttime()
450
+ with torch.no_grad():
451
+ # pred_semantic = t2s_model.model.infer(
452
+ pred_semantic, idx = t2s_model.model.infer_panel(
453
+ all_phoneme_ids,
454
+ all_phoneme_len,
455
+ None if ref_free else prompt,
456
+ bert,
457
+ # prompt_phone_len=ph_offset,
458
+ top_k=top_k,
459
+ top_p=top_p,
460
+ temperature=temperature,
461
+ early_stop_num=hz * max_sec,
462
+ )
463
+ t3 = ttime()
464
+ # print(pred_semantic.shape,idx)
465
+ pred_semantic = pred_semantic[:, -idx:].unsqueeze(
466
+ 0
467
+ ) # .unsqueeze(0)#mq要多unsqueeze一次
468
+ refer = get_spepc(hps, ref_wav_path) # .to(device)
469
+ if is_half == True:
470
+ refer = refer.half().to(device)
471
+ else:
472
+ refer = refer.to(device)
473
+ # audio = vq_model.decode(pred_semantic, all_phoneme_ids, refer).detach().cpu().numpy()[0, 0]
474
+ audio = (
475
+ vq_model.decode(
476
+ pred_semantic, torch.LongTensor(phones2).to(device).unsqueeze(0), refer
477
+ )
478
+ .detach()
479
+ .cpu()
480
+ .numpy()[0, 0]
481
+ ) ###试试重建不带上prompt部分
482
+ max_audio=np.abs(audio).max()#简单防止16bit爆音
483
+ if max_audio>1:audio/=max_audio
484
+ audio_opt.append(audio)
485
+ audio_opt.append(zero_wav)
486
+ t4 = ttime()
487
+ print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
488
+ yield hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(
489
+ np.int16
490
+ )
491
+
492
+
493
+ def split(todo_text):
494
+ todo_text = todo_text.replace("……", "。").replace("——", ",")
495
+ if todo_text[-1] not in splits:
496
+ todo_text += "。"
497
+ i_split_head = i_split_tail = 0
498
+ len_text = len(todo_text)
499
+ todo_texts = []
500
+ while 1:
501
+ if i_split_head >= len_text:
502
+ break # 结尾一定有标点,所以直接跳出即可,最后一段在上次已加入
503
+ if todo_text[i_split_head] in splits:
504
+ i_split_head += 1
505
+ todo_texts.append(todo_text[i_split_tail:i_split_head])
506
+ i_split_tail = i_split_head
507
+ else:
508
+ i_split_head += 1
509
+ return todo_texts
510
+
511
+
512
+ def cut1(inp):
513
+ inp = inp.strip("\n")
514
+ inps = split(inp)
515
+ split_idx = list(range(0, len(inps), 4))
516
+ split_idx[-1] = None
517
+ if len(split_idx) > 1:
518
+ opts = []
519
+ for idx in range(len(split_idx) - 1):
520
+ opts.append("".join(inps[split_idx[idx]: split_idx[idx + 1]]))
521
+ else:
522
+ opts = [inp]
523
+ return "\n".join(opts)
524
+
525
+
526
+ def cut2(inp):
527
+ inp = inp.strip("\n")
528
+ inps = split(inp)
529
+ if len(inps) < 2:
530
+ return inp
531
+ opts = []
532
+ summ = 0
533
+ tmp_str = ""
534
+ for i in range(len(inps)):
535
+ summ += len(inps[i])
536
+ tmp_str += inps[i]
537
+ if summ > 50:
538
+ summ = 0
539
+ opts.append(tmp_str)
540
+ tmp_str = ""
541
+ if tmp_str != "":
542
+ opts.append(tmp_str)
543
+ # print(opts)
544
+ if len(opts) > 1 and len(opts[-1]) < 50: ##如果最后一个太短了,和前一个合一起
545
+ opts[-2] = opts[-2] + opts[-1]
546
+ opts = opts[:-1]
547
+ return "\n".join(opts)
548
+
549
+
550
+ def cut3(inp):
551
+ inp = inp.strip("\n")
552
+ return "\n".join(["%s" % item for item in inp.strip("。").split("。")])
553
+
554
+
555
+ def cut4(inp):
556
+ inp = inp.strip("\n")
557
+ return "\n".join(["%s" % item for item in inp.strip(".").split(".")])
558
+
559
+
560
+ # contributed by https://github.com/AI-Hobbyist/GPT-SoVITS/blob/main/GPT_SoVITS/inference_webui.py
561
+ def cut5(inp):
562
+ # if not re.search(r'[^\w\s]', inp[-1]):
563
+ # inp += '。'
564
+ inp = inp.strip("\n")
565
+ punds = r'[,.;?!、,。?!;:…]'
566
+ items = re.split(f'({punds})', inp)
567
+ mergeitems = ["".join(group) for group in zip(items[::2], items[1::2])]
568
+ # 在句子不存在符号或句尾无符号的时候保证文本完整
569
+ if len(items)%2 == 1:
570
+ mergeitems.append(items[-1])
571
+ opt = "\n".join(mergeitems)
572
+ return opt
573
+
574
+
575
+ def custom_sort_key(s):
576
+ # 使用正则表达式提取字符串中的数字部分和非数字部分
577
+ parts = re.split('(\d+)', s)
578
+ # 将数字部分转换为整数,非数字部分保持不变
579
+ parts = [int(part) if part.isdigit() else part for part in parts]
580
+ return parts
581
+
582
+
583
+ def change_choices():
584
+ SoVITS_names, GPT_names = get_weights_names()
585
+ return {"choices": sorted(SoVITS_names, key=custom_sort_key), "__type__": "update"}, {"choices": sorted(GPT_names, key=custom_sort_key), "__type__": "update"}
586
+
587
+
588
+ pretrained_sovits_name = "GPT_SoVITS/pretrained_models/s2G488k.pth"
589
+ pretrained_gpt_name = "GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"
590
+ SoVITS_weight_root = "SoVITS_weights"
591
+ GPT_weight_root = "GPT_weights"
592
+ os.makedirs(SoVITS_weight_root, exist_ok=True)
593
+ os.makedirs(GPT_weight_root, exist_ok=True)
594
+
595
+
596
+ def get_weights_names():
597
+ SoVITS_names = [pretrained_sovits_name]
598
+ for name in os.listdir(SoVITS_weight_root):
599
+ if name.endswith(".pth"): SoVITS_names.append("%s/%s" % (SoVITS_weight_root, name))
600
+ GPT_names = [pretrained_gpt_name]
601
+ for name in os.listdir(GPT_weight_root):
602
+ if name.endswith(".ckpt"): GPT_names.append("%s/%s" % (GPT_weight_root, name))
603
+ return SoVITS_names, GPT_names
604
+
605
+
606
+ SoVITS_names, GPT_names = get_weights_names()
607
+
608
+ with gr.Blocks(title="GPT-SoVITS WebUI") as app:
609
+ gr.Markdown(
610
+ value=i18n("本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. <br>如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录<b>LICENSE</b>.")
611
+ )
612
+ with gr.Group():
613
+ gr.Markdown(value=i18n("模型切换"))
614
+ with gr.Row():
615
+ GPT_dropdown = gr.Dropdown(label=i18n("GPT模型列表"), choices=sorted(GPT_names, key=custom_sort_key), value=gpt_path, interactive=True)
616
+ SoVITS_dropdown = gr.Dropdown(label=i18n("SoVITS模型列表"), choices=sorted(SoVITS_names, key=custom_sort_key), value=sovits_path, interactive=True)
617
+ refresh_button = gr.Button(i18n("刷新模型路径"), variant="primary")
618
+ refresh_button.click(fn=change_choices, inputs=[], outputs=[SoVITS_dropdown, GPT_dropdown])
619
+ SoVITS_dropdown.change(change_sovits_weights, [SoVITS_dropdown], [])
620
+ GPT_dropdown.change(change_gpt_weights, [GPT_dropdown], [])
621
+ gr.Markdown(value=i18n("*请上传并填写参考信息"))
622
+ with gr.Row():
623
+ inp_ref = gr.Audio(label=i18n("请上传3~10秒内参考音频,超过会报错!"), type="filepath")
624
+ with gr.Column():
625
+ ref_text_free = gr.Checkbox(label=i18n("开启无参考文本模式。不填参考文本亦相当于开启。"), value=False, interactive=True, show_label=True)
626
+ gr.Markdown(i18n("使用无参考文本模式时建议使用微调的GPT"))
627
+ prompt_text = gr.Textbox(label=i18n("参考音频的文本"), value="")
628
+ prompt_language = gr.Dropdown(
629
+ label=i18n("参考音频的语种"), choices=[i18n("中文"), i18n("英文"), i18n("日文"), i18n("中英混合"), i18n("日英混合"), i18n("多语种混合")], value=i18n("中文")
630
+ )
631
+ gr.Markdown(value=i18n("*请填写需要合成的目标文本。中英混合选中文,日英混合选日文,中日混合暂不支持,非目标语言文本自动遗弃。"))
632
+ with gr.Row():
633
+ text = gr.Textbox(label=i18n("需要合成的文本"), value="")
634
+ text_language = gr.Dropdown(
635
+ label=i18n("需要合成的语种"), choices=[i18n("中文"), i18n("英文"), i18n("日文"), i18n("中英混合"), i18n("日英混合"), i18n("多语种混合")], value=i18n("中文")
636
+ )
637
+ how_to_cut = gr.Radio(
638
+ label=i18n("怎么切"),
639
+ choices=[i18n("不切"), i18n("凑四句一切"), i18n("凑50字一切"), i18n("按中文句号。切"), i18n("按英文句号.切"), i18n("按标点符号切"), ],
640
+ value=i18n("凑四句一切"),
641
+ interactive=True,
642
+ )
643
+ with gr.Row():
644
+ gr.Markdown("gpt采样参数(无参考文本时不要太低):")
645
+ top_k = gr.Slider(minimum=1,maximum=100,step=1,label=i18n("top_k"),value=5,interactive=True)
646
+ top_p = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("top_p"),value=1,interactive=True)
647
+ temperature = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("temperature"),value=1,interactive=True)
648
+ inference_button = gr.Button(i18n("合成语音"), variant="primary")
649
+ output = gr.Audio(label=i18n("输出的语音"))
650
+
651
+ inference_button.click(
652
+ get_tts_wav,
653
+ [inp_ref, prompt_text, prompt_language, text, text_language, how_to_cut, top_k, top_p, temperature, ref_text_free],
654
+ [output],
655
+ )
656
+
657
+ gr.Markdown(value=i18n("文本切分工具。太长的文本合成出来效果不一定好,所以太长建议先切。合成会根据文本的换行分开合成再拼起来。"))
658
+ with gr.Row():
659
+ text_inp = gr.Textbox(label=i18n("需要合成的切分前文本"), value="")
660
+ button1 = gr.Button(i18n("凑四句一切"), variant="primary")
661
+ button2 = gr.Button(i18n("凑50字一切"), variant="primary")
662
+ button3 = gr.Button(i18n("按中文句号。切"), variant="primary")
663
+ button4 = gr.Button(i18n("按英文句号.切"), variant="primary")
664
+ button5 = gr.Button(i18n("按标点符号切"), variant="primary")
665
+ text_opt = gr.Textbox(label=i18n("切分后文本"), value="")
666
+ button1.click(cut1, [text_inp], [text_opt])
667
+ button2.click(cut2, [text_inp], [text_opt])
668
+ button3.click(cut3, [text_inp], [text_opt])
669
+ button4.click(cut4, [text_inp], [text_opt])
670
+ button5.click(cut5, [text_inp], [text_opt])
671
+ gr.Markdown(value=i18n("后续将支持转音素、手工修改音素、语音合成分步执行。"))
672
+
673
+ app.queue(concurrency_count=511, max_size=1022).launch(
674
+ server_name="0.0.0.0",
675
+ inbrowser=True,
676
+ share=is_share,
677
+ server_port=infer_ttswebui,
678
+ quiet=True,
679
+ )
lipsync-models/GPT_SoVITS/module/__init__.py ADDED
File without changes
lipsync-models/GPT_SoVITS/module/attentions.py ADDED
@@ -0,0 +1,709 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import torch
3
+ from torch import nn
4
+ from torch.nn import functional as F
5
+
6
+ from module import commons
7
+ from module.modules import LayerNorm
8
+
9
+
10
+ class Encoder(nn.Module):
11
+ def __init__(
12
+ self,
13
+ hidden_channels,
14
+ filter_channels,
15
+ n_heads,
16
+ n_layers,
17
+ kernel_size=1,
18
+ p_dropout=0.0,
19
+ window_size=4,
20
+ isflow=False,
21
+ **kwargs
22
+ ):
23
+ super().__init__()
24
+ self.hidden_channels = hidden_channels
25
+ self.filter_channels = filter_channels
26
+ self.n_heads = n_heads
27
+ self.n_layers = n_layers
28
+ self.kernel_size = kernel_size
29
+ self.p_dropout = p_dropout
30
+ self.window_size = window_size
31
+
32
+ self.drop = nn.Dropout(p_dropout)
33
+ self.attn_layers = nn.ModuleList()
34
+ self.norm_layers_1 = nn.ModuleList()
35
+ self.ffn_layers = nn.ModuleList()
36
+ self.norm_layers_2 = nn.ModuleList()
37
+ for i in range(self.n_layers):
38
+ self.attn_layers.append(
39
+ MultiHeadAttention(
40
+ hidden_channels,
41
+ hidden_channels,
42
+ n_heads,
43
+ p_dropout=p_dropout,
44
+ window_size=window_size,
45
+ )
46
+ )
47
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
48
+ self.ffn_layers.append(
49
+ FFN(
50
+ hidden_channels,
51
+ hidden_channels,
52
+ filter_channels,
53
+ kernel_size,
54
+ p_dropout=p_dropout,
55
+ )
56
+ )
57
+ self.norm_layers_2.append(LayerNorm(hidden_channels))
58
+ if isflow:
59
+ cond_layer = torch.nn.Conv1d(
60
+ kwargs["gin_channels"], 2 * hidden_channels * n_layers, 1
61
+ )
62
+ self.cond_pre = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, 1)
63
+ self.cond_layer = weight_norm_modules(cond_layer, name="weight")
64
+ self.gin_channels = kwargs["gin_channels"]
65
+
66
+ def forward(self, x, x_mask, g=None):
67
+ attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
68
+ x = x * x_mask
69
+ if g is not None:
70
+ g = self.cond_layer(g)
71
+
72
+ for i in range(self.n_layers):
73
+ if g is not None:
74
+ x = self.cond_pre(x)
75
+ cond_offset = i * 2 * self.hidden_channels
76
+ g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
77
+ x = commons.fused_add_tanh_sigmoid_multiply(
78
+ x, g_l, torch.IntTensor([self.hidden_channels])
79
+ )
80
+ y = self.attn_layers[i](x, x, attn_mask)
81
+ y = self.drop(y)
82
+ x = self.norm_layers_1[i](x + y)
83
+
84
+ y = self.ffn_layers[i](x, x_mask)
85
+ y = self.drop(y)
86
+ x = self.norm_layers_2[i](x + y)
87
+ x = x * x_mask
88
+ return x
89
+
90
+
91
+ class Decoder(nn.Module):
92
+ def __init__(
93
+ self,
94
+ hidden_channels,
95
+ filter_channels,
96
+ n_heads,
97
+ n_layers,
98
+ kernel_size=1,
99
+ p_dropout=0.0,
100
+ proximal_bias=False,
101
+ proximal_init=True,
102
+ **kwargs
103
+ ):
104
+ super().__init__()
105
+ self.hidden_channels = hidden_channels
106
+ self.filter_channels = filter_channels
107
+ self.n_heads = n_heads
108
+ self.n_layers = n_layers
109
+ self.kernel_size = kernel_size
110
+ self.p_dropout = p_dropout
111
+ self.proximal_bias = proximal_bias
112
+ self.proximal_init = proximal_init
113
+
114
+ self.drop = nn.Dropout(p_dropout)
115
+ self.self_attn_layers = nn.ModuleList()
116
+ self.norm_layers_0 = nn.ModuleList()
117
+ self.encdec_attn_layers = nn.ModuleList()
118
+ self.norm_layers_1 = nn.ModuleList()
119
+ self.ffn_layers = nn.ModuleList()
120
+ self.norm_layers_2 = nn.ModuleList()
121
+ for i in range(self.n_layers):
122
+ self.self_attn_layers.append(
123
+ MultiHeadAttention(
124
+ hidden_channels,
125
+ hidden_channels,
126
+ n_heads,
127
+ p_dropout=p_dropout,
128
+ proximal_bias=proximal_bias,
129
+ proximal_init=proximal_init,
130
+ )
131
+ )
132
+ self.norm_layers_0.append(LayerNorm(hidden_channels))
133
+ self.encdec_attn_layers.append(
134
+ MultiHeadAttention(
135
+ hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
136
+ )
137
+ )
138
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
139
+ self.ffn_layers.append(
140
+ FFN(
141
+ hidden_channels,
142
+ hidden_channels,
143
+ filter_channels,
144
+ kernel_size,
145
+ p_dropout=p_dropout,
146
+ causal=True,
147
+ )
148
+ )
149
+ self.norm_layers_2.append(LayerNorm(hidden_channels))
150
+
151
+ def forward(self, x, x_mask, h, h_mask):
152
+ """
153
+ x: decoder input
154
+ h: encoder output
155
+ """
156
+ self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
157
+ device=x.device, dtype=x.dtype
158
+ )
159
+ encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
160
+ x = x * x_mask
161
+ for i in range(self.n_layers):
162
+ y = self.self_attn_layers[i](x, x, self_attn_mask)
163
+ y = self.drop(y)
164
+ x = self.norm_layers_0[i](x + y)
165
+
166
+ y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
167
+ y = self.drop(y)
168
+ x = self.norm_layers_1[i](x + y)
169
+
170
+ y = self.ffn_layers[i](x, x_mask)
171
+ y = self.drop(y)
172
+ x = self.norm_layers_2[i](x + y)
173
+ x = x * x_mask
174
+ return x
175
+
176
+
177
+ class MultiHeadAttention(nn.Module):
178
+ def __init__(
179
+ self,
180
+ channels,
181
+ out_channels,
182
+ n_heads,
183
+ p_dropout=0.0,
184
+ window_size=None,
185
+ heads_share=True,
186
+ block_length=None,
187
+ proximal_bias=False,
188
+ proximal_init=False,
189
+ ):
190
+ super().__init__()
191
+ assert channels % n_heads == 0
192
+
193
+ self.channels = channels
194
+ self.out_channels = out_channels
195
+ self.n_heads = n_heads
196
+ self.p_dropout = p_dropout
197
+ self.window_size = window_size
198
+ self.heads_share = heads_share
199
+ self.block_length = block_length
200
+ self.proximal_bias = proximal_bias
201
+ self.proximal_init = proximal_init
202
+ self.attn = None
203
+
204
+ self.k_channels = channels // n_heads
205
+ self.conv_q = nn.Conv1d(channels, channels, 1)
206
+ self.conv_k = nn.Conv1d(channels, channels, 1)
207
+ self.conv_v = nn.Conv1d(channels, channels, 1)
208
+ self.conv_o = nn.Conv1d(channels, out_channels, 1)
209
+ self.drop = nn.Dropout(p_dropout)
210
+
211
+ if window_size is not None:
212
+ n_heads_rel = 1 if heads_share else n_heads
213
+ rel_stddev = self.k_channels**-0.5
214
+ self.emb_rel_k = nn.Parameter(
215
+ torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
216
+ * rel_stddev
217
+ )
218
+ self.emb_rel_v = nn.Parameter(
219
+ torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
220
+ * rel_stddev
221
+ )
222
+
223
+ nn.init.xavier_uniform_(self.conv_q.weight)
224
+ nn.init.xavier_uniform_(self.conv_k.weight)
225
+ nn.init.xavier_uniform_(self.conv_v.weight)
226
+ if proximal_init:
227
+ with torch.no_grad():
228
+ self.conv_k.weight.copy_(self.conv_q.weight)
229
+ self.conv_k.bias.copy_(self.conv_q.bias)
230
+
231
+ def forward(self, x, c, attn_mask=None):
232
+ q = self.conv_q(x)
233
+ k = self.conv_k(c)
234
+ v = self.conv_v(c)
235
+
236
+ x, self.attn = self.attention(q, k, v, mask=attn_mask)
237
+
238
+ x = self.conv_o(x)
239
+ return x
240
+
241
+ def attention(self, query, key, value, mask=None):
242
+ # reshape [b, d, t] -> [b, n_h, t, d_k]
243
+ b, d, t_s, t_t = (*key.size(), query.size(2))
244
+ query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
245
+ key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
246
+ value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
247
+
248
+ scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
249
+ if self.window_size is not None:
250
+ assert (
251
+ t_s == t_t
252
+ ), "Relative attention is only available for self-attention."
253
+ key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
254
+ rel_logits = self._matmul_with_relative_keys(
255
+ query / math.sqrt(self.k_channels), key_relative_embeddings
256
+ )
257
+ scores_local = self._relative_position_to_absolute_position(rel_logits)
258
+ scores = scores + scores_local
259
+ if self.proximal_bias:
260
+ assert t_s == t_t, "Proximal bias is only available for self-attention."
261
+ scores = scores + self._attention_bias_proximal(t_s).to(
262
+ device=scores.device, dtype=scores.dtype
263
+ )
264
+ if mask is not None:
265
+ scores = scores.masked_fill(mask == 0, -1e4)
266
+ if self.block_length is not None:
267
+ assert (
268
+ t_s == t_t
269
+ ), "Local attention is only available for self-attention."
270
+ block_mask = (
271
+ torch.ones_like(scores)
272
+ .triu(-self.block_length)
273
+ .tril(self.block_length)
274
+ )
275
+ scores = scores.masked_fill(block_mask == 0, -1e4)
276
+ p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]
277
+ p_attn = self.drop(p_attn)
278
+ output = torch.matmul(p_attn, value)
279
+ if self.window_size is not None:
280
+ relative_weights = self._absolute_position_to_relative_position(p_attn)
281
+ value_relative_embeddings = self._get_relative_embeddings(
282
+ self.emb_rel_v, t_s
283
+ )
284
+ output = output + self._matmul_with_relative_values(
285
+ relative_weights, value_relative_embeddings
286
+ )
287
+ output = (
288
+ output.transpose(2, 3).contiguous().view(b, d, t_t)
289
+ ) # [b, n_h, t_t, d_k] -> [b, d, t_t]
290
+ return output, p_attn
291
+
292
+ def _matmul_with_relative_values(self, x, y):
293
+ """
294
+ x: [b, h, l, m]
295
+ y: [h or 1, m, d]
296
+ ret: [b, h, l, d]
297
+ """
298
+ ret = torch.matmul(x, y.unsqueeze(0))
299
+ return ret
300
+
301
+ def _matmul_with_relative_keys(self, x, y):
302
+ """
303
+ x: [b, h, l, d]
304
+ y: [h or 1, m, d]
305
+ ret: [b, h, l, m]
306
+ """
307
+ ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
308
+ return ret
309
+
310
+ def _get_relative_embeddings(self, relative_embeddings, length):
311
+ max_relative_position = 2 * self.window_size + 1
312
+ # Pad first before slice to avoid using cond ops.
313
+ pad_length = max(length - (self.window_size + 1), 0)
314
+ slice_start_position = max((self.window_size + 1) - length, 0)
315
+ slice_end_position = slice_start_position + 2 * length - 1
316
+ if pad_length > 0:
317
+ padded_relative_embeddings = F.pad(
318
+ relative_embeddings,
319
+ commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
320
+ )
321
+ else:
322
+ padded_relative_embeddings = relative_embeddings
323
+ used_relative_embeddings = padded_relative_embeddings[
324
+ :, slice_start_position:slice_end_position
325
+ ]
326
+ return used_relative_embeddings
327
+
328
+ def _relative_position_to_absolute_position(self, x):
329
+ """
330
+ x: [b, h, l, 2*l-1]
331
+ ret: [b, h, l, l]
332
+ """
333
+ batch, heads, length, _ = x.size()
334
+ # Concat columns of pad to shift from relative to absolute indexing.
335
+ x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
336
+
337
+ # Concat extra elements so to add up to shape (len+1, 2*len-1).
338
+ x_flat = x.view([batch, heads, length * 2 * length])
339
+ x_flat = F.pad(
340
+ x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
341
+ )
342
+
343
+ # Reshape and slice out the padded elements.
344
+ x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
345
+ :, :, :length, length - 1 :
346
+ ]
347
+ return x_final
348
+
349
+ def _absolute_position_to_relative_position(self, x):
350
+ """
351
+ x: [b, h, l, l]
352
+ ret: [b, h, l, 2*l-1]
353
+ """
354
+ batch, heads, length, _ = x.size()
355
+ # padd along column
356
+ x = F.pad(
357
+ x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
358
+ )
359
+ x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
360
+ # add 0's in the beginning that will skew the elements after reshape
361
+ x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
362
+ x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
363
+ return x_final
364
+
365
+ def _attention_bias_proximal(self, length):
366
+ """Bias for self-attention to encourage attention to close positions.
367
+ Args:
368
+ length: an integer scalar.
369
+ Returns:
370
+ a Tensor with shape [1, 1, length, length]
371
+ """
372
+ r = torch.arange(length, dtype=torch.float32)
373
+ diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
374
+ return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
375
+
376
+
377
+ class FFN(nn.Module):
378
+ def __init__(
379
+ self,
380
+ in_channels,
381
+ out_channels,
382
+ filter_channels,
383
+ kernel_size,
384
+ p_dropout=0.0,
385
+ activation=None,
386
+ causal=False,
387
+ ):
388
+ super().__init__()
389
+ self.in_channels = in_channels
390
+ self.out_channels = out_channels
391
+ self.filter_channels = filter_channels
392
+ self.kernel_size = kernel_size
393
+ self.p_dropout = p_dropout
394
+ self.activation = activation
395
+ self.causal = causal
396
+
397
+ if causal:
398
+ self.padding = self._causal_padding
399
+ else:
400
+ self.padding = self._same_padding
401
+
402
+ self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
403
+ self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
404
+ self.drop = nn.Dropout(p_dropout)
405
+
406
+ def forward(self, x, x_mask):
407
+ x = self.conv_1(self.padding(x * x_mask))
408
+ if self.activation == "gelu":
409
+ x = x * torch.sigmoid(1.702 * x)
410
+ else:
411
+ x = torch.relu(x)
412
+ x = self.drop(x)
413
+ x = self.conv_2(self.padding(x * x_mask))
414
+ return x * x_mask
415
+
416
+ def _causal_padding(self, x):
417
+ if self.kernel_size == 1:
418
+ return x
419
+ pad_l = self.kernel_size - 1
420
+ pad_r = 0
421
+ padding = [[0, 0], [0, 0], [pad_l, pad_r]]
422
+ x = F.pad(x, commons.convert_pad_shape(padding))
423
+ return x
424
+
425
+ def _same_padding(self, x):
426
+ if self.kernel_size == 1:
427
+ return x
428
+ pad_l = (self.kernel_size - 1) // 2
429
+ pad_r = self.kernel_size // 2
430
+ padding = [[0, 0], [0, 0], [pad_l, pad_r]]
431
+ x = F.pad(x, commons.convert_pad_shape(padding))
432
+ return x
433
+
434
+
435
+ import torch.nn as nn
436
+ from torch.nn.utils import remove_weight_norm, weight_norm
437
+
438
+
439
+ class Depthwise_Separable_Conv1D(nn.Module):
440
+ def __init__(
441
+ self,
442
+ in_channels,
443
+ out_channels,
444
+ kernel_size,
445
+ stride=1,
446
+ padding=0,
447
+ dilation=1,
448
+ bias=True,
449
+ padding_mode="zeros", # TODO: refine this type
450
+ device=None,
451
+ dtype=None,
452
+ ):
453
+ super().__init__()
454
+ self.depth_conv = nn.Conv1d(
455
+ in_channels=in_channels,
456
+ out_channels=in_channels,
457
+ kernel_size=kernel_size,
458
+ groups=in_channels,
459
+ stride=stride,
460
+ padding=padding,
461
+ dilation=dilation,
462
+ bias=bias,
463
+ padding_mode=padding_mode,
464
+ device=device,
465
+ dtype=dtype,
466
+ )
467
+ self.point_conv = nn.Conv1d(
468
+ in_channels=in_channels,
469
+ out_channels=out_channels,
470
+ kernel_size=1,
471
+ bias=bias,
472
+ device=device,
473
+ dtype=dtype,
474
+ )
475
+
476
+ def forward(self, input):
477
+ return self.point_conv(self.depth_conv(input))
478
+
479
+ def weight_norm(self):
480
+ self.depth_conv = weight_norm(self.depth_conv, name="weight")
481
+ self.point_conv = weight_norm(self.point_conv, name="weight")
482
+
483
+ def remove_weight_norm(self):
484
+ self.depth_conv = remove_weight_norm(self.depth_conv, name="weight")
485
+ self.point_conv = remove_weight_norm(self.point_conv, name="weight")
486
+
487
+
488
+ class Depthwise_Separable_TransposeConv1D(nn.Module):
489
+ def __init__(
490
+ self,
491
+ in_channels,
492
+ out_channels,
493
+ kernel_size,
494
+ stride=1,
495
+ padding=0,
496
+ output_padding=0,
497
+ bias=True,
498
+ dilation=1,
499
+ padding_mode="zeros", # TODO: refine this type
500
+ device=None,
501
+ dtype=None,
502
+ ):
503
+ super().__init__()
504
+ self.depth_conv = nn.ConvTranspose1d(
505
+ in_channels=in_channels,
506
+ out_channels=in_channels,
507
+ kernel_size=kernel_size,
508
+ groups=in_channels,
509
+ stride=stride,
510
+ output_padding=output_padding,
511
+ padding=padding,
512
+ dilation=dilation,
513
+ bias=bias,
514
+ padding_mode=padding_mode,
515
+ device=device,
516
+ dtype=dtype,
517
+ )
518
+ self.point_conv = nn.Conv1d(
519
+ in_channels=in_channels,
520
+ out_channels=out_channels,
521
+ kernel_size=1,
522
+ bias=bias,
523
+ device=device,
524
+ dtype=dtype,
525
+ )
526
+
527
+ def forward(self, input):
528
+ return self.point_conv(self.depth_conv(input))
529
+
530
+ def weight_norm(self):
531
+ self.depth_conv = weight_norm(self.depth_conv, name="weight")
532
+ self.point_conv = weight_norm(self.point_conv, name="weight")
533
+
534
+ def remove_weight_norm(self):
535
+ remove_weight_norm(self.depth_conv, name="weight")
536
+ remove_weight_norm(self.point_conv, name="weight")
537
+
538
+
539
+ def weight_norm_modules(module, name="weight", dim=0):
540
+ if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
541
+ module, Depthwise_Separable_TransposeConv1D
542
+ ):
543
+ module.weight_norm()
544
+ return module
545
+ else:
546
+ return weight_norm(module, name, dim)
547
+
548
+
549
+ def remove_weight_norm_modules(module, name="weight"):
550
+ if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
551
+ module, Depthwise_Separable_TransposeConv1D
552
+ ):
553
+ module.remove_weight_norm()
554
+ else:
555
+ remove_weight_norm(module, name)
556
+
557
+
558
+ class FFT(nn.Module):
559
+ def __init__(
560
+ self,
561
+ hidden_channels,
562
+ filter_channels,
563
+ n_heads,
564
+ n_layers=1,
565
+ kernel_size=1,
566
+ p_dropout=0.0,
567
+ proximal_bias=False,
568
+ proximal_init=True,
569
+ isflow=False,
570
+ **kwargs
571
+ ):
572
+ super().__init__()
573
+ self.hidden_channels = hidden_channels
574
+ self.filter_channels = filter_channels
575
+ self.n_heads = n_heads
576
+ self.n_layers = n_layers
577
+ self.kernel_size = kernel_size
578
+ self.p_dropout = p_dropout
579
+ self.proximal_bias = proximal_bias
580
+ self.proximal_init = proximal_init
581
+ if isflow:
582
+ cond_layer = torch.nn.Conv1d(
583
+ kwargs["gin_channels"], 2 * hidden_channels * n_layers, 1
584
+ )
585
+ self.cond_pre = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, 1)
586
+ self.cond_layer = weight_norm_modules(cond_layer, name="weight")
587
+ self.gin_channels = kwargs["gin_channels"]
588
+ self.drop = nn.Dropout(p_dropout)
589
+ self.self_attn_layers = nn.ModuleList()
590
+ self.norm_layers_0 = nn.ModuleList()
591
+ self.ffn_layers = nn.ModuleList()
592
+ self.norm_layers_1 = nn.ModuleList()
593
+ for i in range(self.n_layers):
594
+ self.self_attn_layers.append(
595
+ MultiHeadAttention(
596
+ hidden_channels,
597
+ hidden_channels,
598
+ n_heads,
599
+ p_dropout=p_dropout,
600
+ proximal_bias=proximal_bias,
601
+ proximal_init=proximal_init,
602
+ )
603
+ )
604
+ self.norm_layers_0.append(LayerNorm(hidden_channels))
605
+ self.ffn_layers.append(
606
+ FFN(
607
+ hidden_channels,
608
+ hidden_channels,
609
+ filter_channels,
610
+ kernel_size,
611
+ p_dropout=p_dropout,
612
+ causal=True,
613
+ )
614
+ )
615
+ self.norm_layers_1.append(LayerNorm(hidden_channels))
616
+
617
+ def forward(self, x, x_mask, g=None):
618
+ """
619
+ x: decoder input
620
+ h: encoder output
621
+ """
622
+ if g is not None:
623
+ g = self.cond_layer(g)
624
+
625
+ self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
626
+ device=x.device, dtype=x.dtype
627
+ )
628
+ x = x * x_mask
629
+ for i in range(self.n_layers):
630
+ if g is not None:
631
+ x = self.cond_pre(x)
632
+ cond_offset = i * 2 * self.hidden_channels
633
+ g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
634
+ x = commons.fused_add_tanh_sigmoid_multiply(
635
+ x, g_l, torch.IntTensor([self.hidden_channels])
636
+ )
637
+ y = self.self_attn_layers[i](x, x, self_attn_mask)
638
+ y = self.drop(y)
639
+ x = self.norm_layers_0[i](x + y)
640
+
641
+ y = self.ffn_layers[i](x, x_mask)
642
+ y = self.drop(y)
643
+ x = self.norm_layers_1[i](x + y)
644
+ x = x * x_mask
645
+ return x
646
+
647
+
648
+ class TransformerCouplingLayer(nn.Module):
649
+ def __init__(
650
+ self,
651
+ channels,
652
+ hidden_channels,
653
+ kernel_size,
654
+ n_layers,
655
+ n_heads,
656
+ p_dropout=0,
657
+ filter_channels=0,
658
+ mean_only=False,
659
+ wn_sharing_parameter=None,
660
+ gin_channels=0,
661
+ ):
662
+ assert channels % 2 == 0, "channels should be divisible by 2"
663
+ super().__init__()
664
+ self.channels = channels
665
+ self.hidden_channels = hidden_channels
666
+ self.kernel_size = kernel_size
667
+ self.n_layers = n_layers
668
+ self.half_channels = channels // 2
669
+ self.mean_only = mean_only
670
+
671
+ self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
672
+ self.enc = (
673
+ Encoder(
674
+ hidden_channels,
675
+ filter_channels,
676
+ n_heads,
677
+ n_layers,
678
+ kernel_size,
679
+ p_dropout,
680
+ isflow=True,
681
+ gin_channels=gin_channels,
682
+ )
683
+ if wn_sharing_parameter is None
684
+ else wn_sharing_parameter
685
+ )
686
+ self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
687
+ self.post.weight.data.zero_()
688
+ self.post.bias.data.zero_()
689
+
690
+ def forward(self, x, x_mask, g=None, reverse=False):
691
+ x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
692
+ h = self.pre(x0) * x_mask
693
+ h = self.enc(h, x_mask, g=g)
694
+ stats = self.post(h) * x_mask
695
+ if not self.mean_only:
696
+ m, logs = torch.split(stats, [self.half_channels] * 2, 1)
697
+ else:
698
+ m = stats
699
+ logs = torch.zeros_like(m)
700
+
701
+ if not reverse:
702
+ x1 = m + x1 * torch.exp(logs) * x_mask
703
+ x = torch.cat([x0, x1], 1)
704
+ logdet = torch.sum(logs, [1, 2])
705
+ return x, logdet
706
+ else:
707
+ x1 = (x1 - m) * torch.exp(-logs) * x_mask
708
+ x = torch.cat([x0, x1], 1)
709
+ return x
lipsync-models/GPT_SoVITS/module/commons.py ADDED
@@ -0,0 +1,189 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import torch
3
+ from torch.nn import functional as F
4
+
5
+
6
+ def init_weights(m, mean=0.0, std=0.01):
7
+ classname = m.__class__.__name__
8
+ if classname.find("Conv") != -1:
9
+ m.weight.data.normal_(mean, std)
10
+
11
+
12
+ def get_padding(kernel_size, dilation=1):
13
+ return int((kernel_size * dilation - dilation) / 2)
14
+
15
+
16
+ def convert_pad_shape(pad_shape):
17
+ l = pad_shape[::-1]
18
+ pad_shape = [item for sublist in l for item in sublist]
19
+ return pad_shape
20
+
21
+
22
+ def intersperse(lst, item):
23
+ result = [item] * (len(lst) * 2 + 1)
24
+ result[1::2] = lst
25
+ return result
26
+
27
+
28
+ def kl_divergence(m_p, logs_p, m_q, logs_q):
29
+ """KL(P||Q)"""
30
+ kl = (logs_q - logs_p) - 0.5
31
+ kl += (
32
+ 0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
33
+ )
34
+ return kl
35
+
36
+
37
+ def rand_gumbel(shape):
38
+ """Sample from the Gumbel distribution, protect from overflows."""
39
+ uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
40
+ return -torch.log(-torch.log(uniform_samples))
41
+
42
+
43
+ def rand_gumbel_like(x):
44
+ g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
45
+ return g
46
+
47
+
48
+ def slice_segments(x, ids_str, segment_size=4):
49
+ ret = torch.zeros_like(x[:, :, :segment_size])
50
+ for i in range(x.size(0)):
51
+ idx_str = ids_str[i]
52
+ idx_end = idx_str + segment_size
53
+ ret[i] = x[i, :, idx_str:idx_end]
54
+ return ret
55
+
56
+
57
+ def rand_slice_segments(x, x_lengths=None, segment_size=4):
58
+ b, d, t = x.size()
59
+ if x_lengths is None:
60
+ x_lengths = t
61
+ ids_str_max = x_lengths - segment_size + 1
62
+ ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
63
+ ret = slice_segments(x, ids_str, segment_size)
64
+ return ret, ids_str
65
+
66
+
67
+ def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
68
+ position = torch.arange(length, dtype=torch.float)
69
+ num_timescales = channels // 2
70
+ log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
71
+ num_timescales - 1
72
+ )
73
+ inv_timescales = min_timescale * torch.exp(
74
+ torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
75
+ )
76
+ scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
77
+ signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
78
+ signal = F.pad(signal, [0, 0, 0, channels % 2])
79
+ signal = signal.view(1, channels, length)
80
+ return signal
81
+
82
+
83
+ def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
84
+ b, channels, length = x.size()
85
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
86
+ return x + signal.to(dtype=x.dtype, device=x.device)
87
+
88
+
89
+ def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
90
+ b, channels, length = x.size()
91
+ signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
92
+ return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
93
+
94
+
95
+ def subsequent_mask(length):
96
+ mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
97
+ return mask
98
+
99
+
100
+ @torch.jit.script
101
+ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
102
+ n_channels_int = n_channels[0]
103
+ in_act = input_a + input_b
104
+ t_act = torch.tanh(in_act[:, :n_channels_int, :])
105
+ s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
106
+ acts = t_act * s_act
107
+ return acts
108
+
109
+
110
+ def convert_pad_shape(pad_shape):
111
+ l = pad_shape[::-1]
112
+ pad_shape = [item for sublist in l for item in sublist]
113
+ return pad_shape
114
+
115
+
116
+ def shift_1d(x):
117
+ x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
118
+ return x
119
+
120
+
121
+ def sequence_mask(length, max_length=None):
122
+ if max_length is None:
123
+ max_length = length.max()
124
+ x = torch.arange(max_length, dtype=length.dtype, device=length.device)
125
+ return x.unsqueeze(0) < length.unsqueeze(1)
126
+
127
+
128
+ def generate_path(duration, mask):
129
+ """
130
+ duration: [b, 1, t_x]
131
+ mask: [b, 1, t_y, t_x]
132
+ """
133
+ device = duration.device
134
+
135
+ b, _, t_y, t_x = mask.shape
136
+ cum_duration = torch.cumsum(duration, -1)
137
+
138
+ cum_duration_flat = cum_duration.view(b * t_x)
139
+ path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
140
+ path = path.view(b, t_x, t_y)
141
+ path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
142
+ path = path.unsqueeze(1).transpose(2, 3) * mask
143
+ return path
144
+
145
+
146
+ def clip_grad_value_(parameters, clip_value, norm_type=2):
147
+ if isinstance(parameters, torch.Tensor):
148
+ parameters = [parameters]
149
+ parameters = list(filter(lambda p: p.grad is not None, parameters))
150
+ norm_type = float(norm_type)
151
+ if clip_value is not None:
152
+ clip_value = float(clip_value)
153
+
154
+ total_norm = 0
155
+ for p in parameters:
156
+ param_norm = p.grad.data.norm(norm_type)
157
+ total_norm += param_norm.item() ** norm_type
158
+ if clip_value is not None:
159
+ p.grad.data.clamp_(min=-clip_value, max=clip_value)
160
+ total_norm = total_norm ** (1.0 / norm_type)
161
+ return total_norm
162
+
163
+
164
+ def squeeze(x, x_mask=None, n_sqz=2):
165
+ b, c, t = x.size()
166
+
167
+ t = (t // n_sqz) * n_sqz
168
+ x = x[:, :, :t]
169
+ x_sqz = x.view(b, c, t // n_sqz, n_sqz)
170
+ x_sqz = x_sqz.permute(0, 3, 1, 2).contiguous().view(b, c * n_sqz, t // n_sqz)
171
+
172
+ if x_mask is not None:
173
+ x_mask = x_mask[:, :, n_sqz - 1 :: n_sqz]
174
+ else:
175
+ x_mask = torch.ones(b, 1, t // n_sqz).to(device=x.device, dtype=x.dtype)
176
+ return x_sqz * x_mask, x_mask
177
+
178
+
179
+ def unsqueeze(x, x_mask=None, n_sqz=2):
180
+ b, c, t = x.size()
181
+
182
+ x_unsqz = x.view(b, n_sqz, c // n_sqz, t)
183
+ x_unsqz = x_unsqz.permute(0, 2, 3, 1).contiguous().view(b, c // n_sqz, t * n_sqz)
184
+
185
+ if x_mask is not None:
186
+ x_mask = x_mask.unsqueeze(-1).repeat(1, 1, 1, n_sqz).view(b, 1, t * n_sqz)
187
+ else:
188
+ x_mask = torch.ones(b, 1, t * n_sqz).to(device=x.device, dtype=x.dtype)
189
+ return x_unsqz * x_mask, x_mask
lipsync-models/GPT_SoVITS/module/core_vq.py ADDED
@@ -0,0 +1,383 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+ #
7
+ # This implementation is inspired from
8
+ # https://github.com/lucidrains/vector-quantize-pytorch
9
+ # which is released under MIT License. Hereafter, the original license:
10
+ # MIT License
11
+ #
12
+ # Copyright (c) 2020 Phil Wang
13
+ #
14
+ # Permission is hereby granted, free of charge, to any person obtaining a copy
15
+ # of this software and associated documentation files (the "Software"), to deal
16
+ # in the Software without restriction, including without limitation the rights
17
+ # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
18
+ # copies of the Software, and to permit persons to whom the Software is
19
+ # furnished to do so, subject to the following conditions:
20
+ #
21
+ # The above copyright notice and this permission notice shall be included in all
22
+ # copies or substantial portions of the Software.
23
+ #
24
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
25
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
26
+ # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
27
+ # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
28
+ # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
29
+ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
30
+ # SOFTWARE.
31
+
32
+ """Core vector quantization implementation."""
33
+ import typing as tp
34
+
35
+ from einops import rearrange, repeat
36
+ import torch
37
+ from torch import nn
38
+ import torch.nn.functional as F
39
+ from tqdm import tqdm
40
+
41
+
42
+ def default(val: tp.Any, d: tp.Any) -> tp.Any:
43
+ return val if val is not None else d
44
+
45
+
46
+ def ema_inplace(moving_avg, new, decay: float):
47
+ moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
48
+
49
+
50
+ def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
51
+ return (x + epsilon) / (x.sum() + n_categories * epsilon)
52
+
53
+
54
+ def uniform_init(*shape: int):
55
+ t = torch.empty(shape)
56
+ nn.init.kaiming_uniform_(t)
57
+ return t
58
+
59
+
60
+ def sample_vectors(samples, num: int):
61
+ num_samples, device = samples.shape[0], samples.device
62
+
63
+ if num_samples >= num:
64
+ indices = torch.randperm(num_samples, device=device)[:num]
65
+ else:
66
+ indices = torch.randint(0, num_samples, (num,), device=device)
67
+
68
+ return samples[indices]
69
+
70
+
71
+ def kmeans(samples, num_clusters: int, num_iters: int = 10):
72
+ dim, dtype = samples.shape[-1], samples.dtype
73
+ max_kmeans_samples = 500
74
+ samples = samples[:max_kmeans_samples, :]
75
+ means = sample_vectors(samples, num_clusters)
76
+
77
+ print("kmeans start ... ")
78
+ for _ in tqdm(range(num_iters)):
79
+ diffs = rearrange(samples, "n d -> n () d") - rearrange(means, "c d -> () c d")
80
+ dists = -(diffs**2).sum(dim=-1)
81
+
82
+ buckets = dists.max(dim=-1).indices
83
+ bins = torch.bincount(buckets, minlength=num_clusters)
84
+ zero_mask = bins == 0
85
+ bins_min_clamped = bins.masked_fill(zero_mask, 1)
86
+
87
+ new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
88
+ new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
89
+ new_means = new_means / bins_min_clamped[..., None]
90
+
91
+ means = torch.where(zero_mask[..., None], means, new_means)
92
+
93
+ return means, bins
94
+
95
+
96
+ class EuclideanCodebook(nn.Module):
97
+ """Codebook with Euclidean distance.
98
+ Args:
99
+ dim (int): Dimension.
100
+ codebook_size (int): Codebook size.
101
+ kmeans_init (bool): Whether to use k-means to initialize the codebooks.
102
+ If set to true, run the k-means algorithm on the first training batch and use
103
+ the learned centroids as initialization.
104
+ kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
105
+ decay (float): Decay for exponential moving average over the codebooks.
106
+ epsilon (float): Epsilon value for numerical stability.
107
+ threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
108
+ that have an exponential moving average cluster size less than the specified threshold with
109
+ randomly selected vector from the current batch.
110
+ """
111
+
112
+ def __init__(
113
+ self,
114
+ dim: int,
115
+ codebook_size: int,
116
+ kmeans_init: int = False,
117
+ kmeans_iters: int = 10,
118
+ decay: float = 0.99,
119
+ epsilon: float = 1e-5,
120
+ threshold_ema_dead_code: int = 2,
121
+ ):
122
+ super().__init__()
123
+ self.decay = decay
124
+ init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = (
125
+ uniform_init if not kmeans_init else torch.zeros
126
+ )
127
+ embed = init_fn(codebook_size, dim)
128
+
129
+ self.codebook_size = codebook_size
130
+
131
+ self.kmeans_iters = kmeans_iters
132
+ self.epsilon = epsilon
133
+ self.threshold_ema_dead_code = threshold_ema_dead_code
134
+
135
+ self.register_buffer("inited", torch.Tensor([not kmeans_init]))
136
+ self.register_buffer("cluster_size", torch.zeros(codebook_size))
137
+ self.register_buffer("embed", embed)
138
+ self.register_buffer("embed_avg", embed.clone())
139
+
140
+ @torch.jit.ignore
141
+ def init_embed_(self, data):
142
+ if self.inited:
143
+ return
144
+
145
+ embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
146
+ self.embed.data.copy_(embed)
147
+ self.embed_avg.data.copy_(embed.clone())
148
+ self.cluster_size.data.copy_(cluster_size)
149
+ self.inited.data.copy_(torch.Tensor([True]))
150
+ # Make sure all buffers across workers are in sync after initialization
151
+ # broadcast_tensors(self.buffers())
152
+
153
+ def replace_(self, samples, mask):
154
+ modified_codebook = torch.where(
155
+ mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
156
+ )
157
+ self.embed.data.copy_(modified_codebook)
158
+
159
+ def expire_codes_(self, batch_samples):
160
+ if self.threshold_ema_dead_code == 0:
161
+ return
162
+
163
+ expired_codes = self.cluster_size < self.threshold_ema_dead_code
164
+ if not torch.any(expired_codes):
165
+ return
166
+
167
+ batch_samples = rearrange(batch_samples, "... d -> (...) d")
168
+ self.replace_(batch_samples, mask=expired_codes)
169
+ # broadcast_tensors(self.buffers())
170
+
171
+ def preprocess(self, x):
172
+ x = rearrange(x, "... d -> (...) d")
173
+ return x
174
+
175
+ def quantize(self, x):
176
+ embed = self.embed.t()
177
+ dist = -(
178
+ x.pow(2).sum(1, keepdim=True)
179
+ - 2 * x @ embed
180
+ + embed.pow(2).sum(0, keepdim=True)
181
+ )
182
+ embed_ind = dist.max(dim=-1).indices
183
+ return embed_ind
184
+
185
+ def postprocess_emb(self, embed_ind, shape):
186
+ return embed_ind.view(*shape[:-1])
187
+
188
+ def dequantize(self, embed_ind):
189
+ quantize = F.embedding(embed_ind, self.embed)
190
+ return quantize
191
+
192
+ def encode(self, x):
193
+ shape = x.shape
194
+ # pre-process
195
+ x = self.preprocess(x)
196
+ # quantize
197
+ embed_ind = self.quantize(x)
198
+ # post-process
199
+ embed_ind = self.postprocess_emb(embed_ind, shape)
200
+ return embed_ind
201
+
202
+ def decode(self, embed_ind):
203
+ quantize = self.dequantize(embed_ind)
204
+ return quantize
205
+
206
+ def forward(self, x):
207
+ shape, dtype = x.shape, x.dtype
208
+ x = self.preprocess(x)
209
+
210
+ self.init_embed_(x)
211
+
212
+ embed_ind = self.quantize(x)
213
+ embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
214
+ embed_ind = self.postprocess_emb(embed_ind, shape)
215
+ quantize = self.dequantize(embed_ind)
216
+
217
+ if self.training:
218
+ # We do the expiry of code at that point as buffers are in sync
219
+ # and all the workers will take the same decision.
220
+ self.expire_codes_(x)
221
+ ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
222
+ embed_sum = x.t() @ embed_onehot
223
+ ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
224
+ cluster_size = (
225
+ laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
226
+ * self.cluster_size.sum()
227
+ )
228
+ embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
229
+ self.embed.data.copy_(embed_normalized)
230
+
231
+ return quantize, embed_ind
232
+
233
+
234
+ class VectorQuantization(nn.Module):
235
+ """Vector quantization implementation.
236
+ Currently supports only euclidean distance.
237
+ Args:
238
+ dim (int): Dimension
239
+ codebook_size (int): Codebook size
240
+ codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
241
+ decay (float): Decay for exponential moving average over the codebooks.
242
+ epsilon (float): Epsilon value for numerical stability.
243
+ kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
244
+ kmeans_iters (int): Number of iterations used for kmeans initialization.
245
+ threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
246
+ that have an exponential moving average cluster size less than the specified threshold with
247
+ randomly selected vector from the current batch.
248
+ commitment_weight (float): Weight for commitment loss.
249
+ """
250
+
251
+ def __init__(
252
+ self,
253
+ dim: int,
254
+ codebook_size: int,
255
+ codebook_dim: tp.Optional[int] = None,
256
+ decay: float = 0.99,
257
+ epsilon: float = 1e-5,
258
+ kmeans_init: bool = True,
259
+ kmeans_iters: int = 50,
260
+ threshold_ema_dead_code: int = 2,
261
+ commitment_weight: float = 1.0,
262
+ ):
263
+ super().__init__()
264
+ _codebook_dim: int = default(codebook_dim, dim)
265
+
266
+ requires_projection = _codebook_dim != dim
267
+ self.project_in = (
268
+ nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()
269
+ )
270
+ self.project_out = (
271
+ nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()
272
+ )
273
+
274
+ self.epsilon = epsilon
275
+ self.commitment_weight = commitment_weight
276
+
277
+ self._codebook = EuclideanCodebook(
278
+ dim=_codebook_dim,
279
+ codebook_size=codebook_size,
280
+ kmeans_init=kmeans_init,
281
+ kmeans_iters=kmeans_iters,
282
+ decay=decay,
283
+ epsilon=epsilon,
284
+ threshold_ema_dead_code=threshold_ema_dead_code,
285
+ )
286
+ self.codebook_size = codebook_size
287
+
288
+ @property
289
+ def codebook(self):
290
+ return self._codebook.embed
291
+
292
+ def encode(self, x):
293
+ x = rearrange(x, "b d n -> b n d")
294
+ x = self.project_in(x)
295
+ embed_in = self._codebook.encode(x)
296
+ return embed_in
297
+
298
+ def decode(self, embed_ind):
299
+ quantize = self._codebook.decode(embed_ind)
300
+ quantize = self.project_out(quantize)
301
+ quantize = rearrange(quantize, "b n d -> b d n")
302
+ return quantize
303
+
304
+ def forward(self, x):
305
+ device = x.device
306
+ x = rearrange(x, "b d n -> b n d")
307
+ x = self.project_in(x)
308
+
309
+ quantize, embed_ind = self._codebook(x)
310
+
311
+ if self.training:
312
+ quantize = x + (quantize - x).detach()
313
+
314
+ loss = torch.tensor([0.0], device=device, requires_grad=self.training)
315
+
316
+ if self.training:
317
+ if self.commitment_weight > 0:
318
+ commit_loss = F.mse_loss(quantize.detach(), x)
319
+ loss = loss + commit_loss * self.commitment_weight
320
+
321
+ quantize = self.project_out(quantize)
322
+ quantize = rearrange(quantize, "b n d -> b d n")
323
+ return quantize, embed_ind, loss
324
+
325
+
326
+ class ResidualVectorQuantization(nn.Module):
327
+ """Residual vector quantization implementation.
328
+ Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
329
+ """
330
+
331
+ def __init__(self, *, num_quantizers, **kwargs):
332
+ super().__init__()
333
+ self.layers = nn.ModuleList(
334
+ [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
335
+ )
336
+
337
+ def forward(
338
+ self, x, n_q: tp.Optional[int] = None, layers: tp.Optional[list] = None
339
+ ):
340
+ quantized_out = 0.0
341
+ residual = x
342
+
343
+ all_losses = []
344
+ all_indices = []
345
+ out_quantized = []
346
+
347
+ n_q = n_q or len(self.layers)
348
+
349
+ for i, layer in enumerate(self.layers[:n_q]):
350
+ quantized, indices, loss = layer(residual)
351
+ residual = residual - quantized
352
+ quantized_out = quantized_out + quantized
353
+
354
+ all_indices.append(indices)
355
+ all_losses.append(loss)
356
+ if layers and i in layers:
357
+ out_quantized.append(quantized)
358
+
359
+ out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
360
+ return quantized_out, out_indices, out_losses, out_quantized
361
+
362
+ def encode(
363
+ self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None
364
+ ) -> torch.Tensor:
365
+ residual = x
366
+ all_indices = []
367
+ n_q = n_q or len(self.layers)
368
+ st = st or 0
369
+ for layer in self.layers[st:n_q]:
370
+ indices = layer.encode(residual)
371
+ quantized = layer.decode(indices)
372
+ residual = residual - quantized
373
+ all_indices.append(indices)
374
+ out_indices = torch.stack(all_indices)
375
+ return out_indices
376
+
377
+ def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor:
378
+ quantized_out = torch.tensor(0.0, device=q_indices.device)
379
+ for i, indices in enumerate(q_indices):
380
+ layer = self.layers[st + i]
381
+ quantized = layer.decode(indices)
382
+ quantized_out = quantized_out + quantized
383
+ return quantized_out
lipsync-models/GPT_SoVITS/module/data_utils.py ADDED
@@ -0,0 +1,332 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ import logging
3
+ import os
4
+ import random
5
+ import traceback
6
+ import numpy as np
7
+ import torch
8
+ import torch.utils.data
9
+ from tqdm import tqdm
10
+
11
+ from module import commons
12
+ from module.mel_processing import spectrogram_torch
13
+ from text import cleaned_text_to_sequence
14
+ from utils import load_wav_to_torch, load_filepaths_and_text
15
+ import torch.nn.functional as F
16
+ from functools import lru_cache
17
+ import requests
18
+ from scipy.io import wavfile
19
+ from io import BytesIO
20
+ from my_utils import load_audio
21
+
22
+ # ZeroDivisionError fixed by Tybost (https://github.com/RVC-Boss/GPT-SoVITS/issues/79)
23
+ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
24
+ """
25
+ 1) loads audio, speaker_id, text pairs
26
+ 2) normalizes text and converts them to sequences of integers
27
+ 3) computes spectrograms from audio files.
28
+ """
29
+
30
+ def __init__(self, hparams, val=False):
31
+ exp_dir = hparams.exp_dir
32
+ self.path2 = "%s/2-name2text.txt" % exp_dir
33
+ self.path4 = "%s/4-cnhubert" % exp_dir
34
+ self.path5 = "%s/5-wav32k" % exp_dir
35
+ assert os.path.exists(self.path2)
36
+ assert os.path.exists(self.path4)
37
+ assert os.path.exists(self.path5)
38
+ names4 = set([name[:-3] for name in list(os.listdir(self.path4))]) # 去除.pt后缀
39
+ names5 = set(os.listdir(self.path5))
40
+ self.phoneme_data = {}
41
+ with open(self.path2, "r", encoding="utf8") as f:
42
+ lines = f.read().strip("\n").split("\n")
43
+
44
+ for line in lines:
45
+ tmp = line.split("\t")
46
+ if (len(tmp) != 4):
47
+ continue
48
+ self.phoneme_data[tmp[0]] = [tmp[1]]
49
+
50
+ self.audiopaths_sid_text = list(set(self.phoneme_data) & names4 & names5)
51
+ tmp = self.audiopaths_sid_text
52
+ leng = len(tmp)
53
+ min_num = 100
54
+ if (leng < min_num):
55
+ self.audiopaths_sid_text = []
56
+ for _ in range(max(2, int(min_num / leng))):
57
+ self.audiopaths_sid_text += tmp
58
+ self.max_wav_value = hparams.max_wav_value
59
+ self.sampling_rate = hparams.sampling_rate
60
+ self.filter_length = hparams.filter_length
61
+ self.hop_length = hparams.hop_length
62
+ self.win_length = hparams.win_length
63
+ self.sampling_rate = hparams.sampling_rate
64
+ self.val = val
65
+
66
+ random.seed(1234)
67
+ random.shuffle(self.audiopaths_sid_text)
68
+
69
+ print("phoneme_data_len:", len(self.phoneme_data.keys()))
70
+ print("wav_data_len:", len(self.audiopaths_sid_text))
71
+
72
+ audiopaths_sid_text_new = []
73
+ lengths = []
74
+ skipped_phone = 0
75
+ skipped_dur = 0
76
+ for audiopath in tqdm(self.audiopaths_sid_text):
77
+ try:
78
+ phoneme = self.phoneme_data[audiopath][0]
79
+ phoneme = phoneme.split(' ')
80
+ phoneme_ids = cleaned_text_to_sequence(phoneme)
81
+ except Exception:
82
+ print(f"{audiopath} not in self.phoneme_data !")
83
+ skipped_phone += 1
84
+ continue
85
+
86
+ size = os.path.getsize("%s/%s" % (self.path5, audiopath))
87
+ duration = size / self.sampling_rate / 2
88
+
89
+ if duration == 0:
90
+ print(f"Zero duration for {audiopath}, skipping...")
91
+ skipped_dur += 1
92
+ continue
93
+
94
+ if 54 > duration > 0.6 or self.val:
95
+ audiopaths_sid_text_new.append([audiopath, phoneme_ids])
96
+ lengths.append(size // (2 * self.hop_length))
97
+ else:
98
+ skipped_dur += 1
99
+ continue
100
+
101
+ print("skipped_phone: ", skipped_phone, ", skipped_dur: ", skipped_dur)
102
+ print("total left: ", len(audiopaths_sid_text_new))
103
+ assert len(audiopaths_sid_text_new) > 1 # 至少能凑够batch size,这里todo
104
+ self.audiopaths_sid_text = audiopaths_sid_text_new
105
+ self.lengths = lengths
106
+
107
+ def get_audio_text_speaker_pair(self, audiopath_sid_text):
108
+ audiopath, phoneme_ids = audiopath_sid_text
109
+ text = torch.FloatTensor(phoneme_ids)
110
+ try:
111
+ spec, wav = self.get_audio("%s/%s" % (self.path5, audiopath))
112
+ with torch.no_grad():
113
+ ssl = torch.load("%s/%s.pt" % (self.path4, audiopath), map_location="cpu")
114
+ if (ssl.shape[-1] != spec.shape[-1]):
115
+ typee = ssl.dtype
116
+ ssl = F.pad(ssl.float(), (0, 1), mode="replicate").to(typee)
117
+ ssl.requires_grad = False
118
+ except:
119
+ traceback.print_exc()
120
+ spec = torch.zeros(1025, 100)
121
+ wav = torch.zeros(1, 100 * self.hop_length)
122
+ ssl = torch.zeros(1, 768, 100)
123
+ text = text[-1:]
124
+ print("load audio or ssl error!!!!!!", audiopath)
125
+ return (ssl, spec, wav, text)
126
+
127
+ def get_audio(self, filename):
128
+ audio_array = load_audio(filename, self.sampling_rate) # load_audio的方法是已经归一化到-1~1之间的,不用再/32768
129
+ audio = torch.FloatTensor(audio_array) # /32768
130
+ audio_norm = audio
131
+ audio_norm = audio_norm.unsqueeze(0)
132
+ spec = spectrogram_torch(audio_norm, self.filter_length, self.sampling_rate, self.hop_length, self.win_length,
133
+ center=False)
134
+ spec = torch.squeeze(spec, 0)
135
+ return spec, audio_norm
136
+
137
+ def get_sid(self, sid):
138
+ sid = torch.LongTensor([int(sid)])
139
+ return sid
140
+
141
+ def __getitem__(self, index):
142
+ # with torch.no_grad():
143
+ return self.get_audio_text_speaker_pair(self.audiopaths_sid_text[index])
144
+
145
+ def __len__(self):
146
+ return len(self.audiopaths_sid_text)
147
+
148
+ def random_slice(self, ssl, wav, mel):
149
+ assert abs(ssl.shape[-1] - wav.shape[-1] // self.hop_length) < 3, (
150
+ "first", ssl.shape, wav.shape)
151
+
152
+ len_mel = mel.shape[1]
153
+ if self.val:
154
+ reference_mel = mel[:, :len_mel // 3]
155
+ return reference_mel, ssl, wav, mel
156
+ dir = random.randint(0, 1)
157
+ sep_point = random.randint(int(len_mel // 3), int(len_mel // 3 * 2))
158
+
159
+ if dir == 0:
160
+ reference_mel = mel[:, :sep_point]
161
+ ssl = ssl[:, :, sep_point:]
162
+ wav2 = wav[:, sep_point * self.hop_length:]
163
+ mel = mel[:, sep_point:]
164
+ else:
165
+ reference_mel = mel[:, sep_point:]
166
+ ssl = ssl[:, :, :sep_point]
167
+ wav2 = wav[:, :sep_point * self.hop_length]
168
+ mel = mel[:, :sep_point]
169
+
170
+ assert abs(ssl.shape[-1] - wav2.shape[-1] // self.hop_length) < 3, (
171
+ ssl.shape, wav.shape, wav2.shape, mel.shape, sep_point, self.hop_length, sep_point * self.hop_length, dir)
172
+ return reference_mel, ssl, wav2, mel
173
+
174
+
175
+ class TextAudioSpeakerCollate():
176
+ """ Zero-pads model inputs and targets
177
+ """
178
+
179
+ def __init__(self, return_ids=False):
180
+ self.return_ids = return_ids
181
+
182
+ def __call__(self, batch):
183
+ """Collate's training batch from normalized text, audio and speaker identities
184
+ PARAMS
185
+ ------
186
+ batch: [text_normalized, spec_normalized, wav_normalized, sid]
187
+ """
188
+ # Right zero-pad all one-hot text sequences to max input length
189
+ _, ids_sorted_decreasing = torch.sort(
190
+ torch.LongTensor([x[1].size(1) for x in batch]),
191
+ dim=0, descending=True)
192
+
193
+ max_ssl_len = max([x[0].size(2) for x in batch])
194
+ max_ssl_len = int(2 * ((max_ssl_len // 2) + 1))
195
+ max_spec_len = max([x[1].size(1) for x in batch])
196
+ max_spec_len = int(2 * ((max_spec_len // 2) + 1))
197
+ max_wav_len = max([x[2].size(1) for x in batch])
198
+ max_text_len = max([x[3].size(0) for x in batch])
199
+
200
+ ssl_lengths = torch.LongTensor(len(batch))
201
+ spec_lengths = torch.LongTensor(len(batch))
202
+ wav_lengths = torch.LongTensor(len(batch))
203
+ text_lengths = torch.LongTensor(len(batch))
204
+
205
+ spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
206
+ wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
207
+ ssl_padded = torch.FloatTensor(len(batch), batch[0][0].size(1), max_ssl_len)
208
+ text_padded = torch.LongTensor(len(batch), max_text_len)
209
+
210
+ spec_padded.zero_()
211
+ wav_padded.zero_()
212
+ ssl_padded.zero_()
213
+ text_padded.zero_()
214
+
215
+ for i in range(len(ids_sorted_decreasing)):
216
+ row = batch[ids_sorted_decreasing[i]]
217
+
218
+ ssl = row[0]
219
+ ssl_padded[i, :, :ssl.size(2)] = ssl[0, :, :]
220
+ ssl_lengths[i] = ssl.size(2)
221
+
222
+ spec = row[1]
223
+ spec_padded[i, :, :spec.size(1)] = spec
224
+ spec_lengths[i] = spec.size(1)
225
+
226
+ wav = row[2]
227
+ wav_padded[i, :, :wav.size(1)] = wav
228
+ wav_lengths[i] = wav.size(1)
229
+
230
+ text = row[3]
231
+ text_padded[i, :text.size(0)] = text
232
+ text_lengths[i] = text.size(0)
233
+
234
+ return ssl_padded, ssl_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, text_padded, text_lengths
235
+
236
+
237
+ class DistributedBucketSampler(torch.utils.data.distributed.DistributedSampler):
238
+ """
239
+ Maintain similar input lengths in a batch.
240
+ Length groups are specified by boundaries.
241
+ Ex) boundaries = [b1, b2, b3] -> any batch is included either {x | b1 < length(x) <=b2} or {x | b2 < length(x) <= b3}.
242
+
243
+ It removes samples which are not included in the boundaries.
244
+ Ex) boundaries = [b1, b2, b3] -> any x s.t. length(x) <= b1 or length(x) > b3 are discarded.
245
+ """
246
+
247
+ def __init__(self, dataset, batch_size, boundaries, num_replicas=None, rank=None, shuffle=True):
248
+ super().__init__(dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle)
249
+ self.lengths = dataset.lengths
250
+ self.batch_size = batch_size
251
+ self.boundaries = boundaries
252
+
253
+ self.buckets, self.num_samples_per_bucket = self._create_buckets()
254
+ self.total_size = sum(self.num_samples_per_bucket)
255
+ self.num_samples = self.total_size // self.num_replicas
256
+
257
+ def _create_buckets(self):
258
+ buckets = [[] for _ in range(len(self.boundaries) - 1)]
259
+ for i in range(len(self.lengths)):
260
+ length = self.lengths[i]
261
+ idx_bucket = self._bisect(length)
262
+ if idx_bucket != -1:
263
+ buckets[idx_bucket].append(i)
264
+
265
+ i = len(buckets) - 1
266
+ while i >= 0:
267
+ if len(buckets[i]) == 0:
268
+ buckets.pop(i)
269
+ self.boundaries.pop(i + 1)
270
+ i -= 1
271
+
272
+ num_samples_per_bucket = []
273
+ for i in range(len(buckets)):
274
+ len_bucket = len(buckets[i])
275
+ total_batch_size = self.num_replicas * self.batch_size
276
+ rem = (total_batch_size - (len_bucket % total_batch_size)) % total_batch_size
277
+ num_samples_per_bucket.append(len_bucket + rem)
278
+ return buckets, num_samples_per_bucket
279
+
280
+ def __iter__(self):
281
+ g = torch.Generator()
282
+ g.manual_seed(self.epoch)
283
+
284
+ indices = []
285
+ if self.shuffle:
286
+ for bucket in self.buckets:
287
+ indices.append(torch.randperm(len(bucket), generator=g).tolist())
288
+ else:
289
+ for bucket in self.buckets:
290
+ indices.append(list(range(len(bucket))))
291
+
292
+ batches = []
293
+ for i in range(len(self.buckets)):
294
+ bucket = self.buckets[i]
295
+ len_bucket = len(bucket)
296
+ ids_bucket = indices[i]
297
+ num_samples_bucket = self.num_samples_per_bucket[i]
298
+
299
+ rem = num_samples_bucket - len_bucket
300
+ ids_bucket = ids_bucket + ids_bucket * (rem // len_bucket) + ids_bucket[:(rem % len_bucket)]
301
+
302
+ ids_bucket = ids_bucket[self.rank::self.num_replicas]
303
+
304
+ for j in range(len(ids_bucket) // self.batch_size):
305
+ batch = [bucket[idx] for idx in ids_bucket[j * self.batch_size:(j + 1) * self.batch_size]]
306
+ batches.append(batch)
307
+
308
+ if self.shuffle:
309
+ batch_ids = torch.randperm(len(batches), generator=g).tolist()
310
+ batches = [batches[i] for i in batch_ids]
311
+ self.batches = batches
312
+
313
+ assert len(self.batches) * self.batch_size == self.num_samples
314
+ return iter(self.batches)
315
+
316
+ def _bisect(self, x, lo=0, hi=None):
317
+ if hi is None:
318
+ hi = len(self.boundaries) - 1
319
+
320
+ if hi > lo:
321
+ mid = (hi + lo) // 2
322
+ if self.boundaries[mid] < x and x <= self.boundaries[mid + 1]:
323
+ return mid
324
+ elif x <= self.boundaries[mid]:
325
+ return self._bisect(x, lo, mid)
326
+ else:
327
+ return self._bisect(x, mid + 1, hi)
328
+ else:
329
+ return -1
330
+
331
+ def __len__(self):
332
+ return self.num_samples // self.batch_size
lipsync-models/GPT_SoVITS/module/losses.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ from torch.nn import functional as F
5
+
6
+
7
+ def feature_loss(fmap_r, fmap_g):
8
+ loss = 0
9
+ for dr, dg in zip(fmap_r, fmap_g):
10
+ for rl, gl in zip(dr, dg):
11
+ rl = rl.float().detach()
12
+ gl = gl.float()
13
+ loss += torch.mean(torch.abs(rl - gl))
14
+
15
+ return loss * 2
16
+
17
+
18
+ def discriminator_loss(disc_real_outputs, disc_generated_outputs):
19
+ loss = 0
20
+ r_losses = []
21
+ g_losses = []
22
+ for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
23
+ dr = dr.float()
24
+ dg = dg.float()
25
+ r_loss = torch.mean((1 - dr) ** 2)
26
+ g_loss = torch.mean(dg**2)
27
+ loss += r_loss + g_loss
28
+ r_losses.append(r_loss.item())
29
+ g_losses.append(g_loss.item())
30
+
31
+ return loss, r_losses, g_losses
32
+
33
+
34
+ def generator_loss(disc_outputs):
35
+ loss = 0
36
+ gen_losses = []
37
+ for dg in disc_outputs:
38
+ dg = dg.float()
39
+ l = torch.mean((1 - dg) ** 2)
40
+ gen_losses.append(l)
41
+ loss += l
42
+
43
+ return loss, gen_losses
44
+
45
+
46
+ def kl_loss(z_p, logs_q, m_p, logs_p, z_mask):
47
+ """
48
+ z_p, logs_q: [b, h, t_t]
49
+ m_p, logs_p: [b, h, t_t]
50
+ """
51
+ z_p = z_p.float()
52
+ logs_q = logs_q.float()
53
+ m_p = m_p.float()
54
+ logs_p = logs_p.float()
55
+ z_mask = z_mask.float()
56
+
57
+ kl = logs_p - logs_q - 0.5
58
+ kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p)
59
+ kl = torch.sum(kl * z_mask)
60
+ l = kl / torch.sum(z_mask)
61
+ return l
62
+
63
+
64
+ def mle_loss(z, m, logs, logdet, mask):
65
+ l = torch.sum(logs) + 0.5 * torch.sum(
66
+ torch.exp(-2 * logs) * ((z - m) ** 2)
67
+ ) # neg normal likelihood w/o the constant term
68
+ l = l - torch.sum(logdet) # log jacobian determinant
69
+ l = l / torch.sum(
70
+ torch.ones_like(z) * mask
71
+ ) # averaging across batch, channel and time axes
72
+ l = l + 0.5 * math.log(2 * math.pi) # add the remaining constant term
73
+ return l
lipsync-models/GPT_SoVITS/module/mel_processing.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import os
3
+ import random
4
+ import torch
5
+ from torch import nn
6
+ import torch.nn.functional as F
7
+ import torch.utils.data
8
+ import numpy as np
9
+ import librosa
10
+ import librosa.util as librosa_util
11
+ from librosa.util import normalize, pad_center, tiny
12
+ from scipy.signal import get_window
13
+ from scipy.io.wavfile import read
14
+ from librosa.filters import mel as librosa_mel_fn
15
+
16
+ MAX_WAV_VALUE = 32768.0
17
+
18
+
19
+ def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
20
+ """
21
+ PARAMS
22
+ ------
23
+ C: compression factor
24
+ """
25
+ return torch.log(torch.clamp(x, min=clip_val) * C)
26
+
27
+
28
+ def dynamic_range_decompression_torch(x, C=1):
29
+ """
30
+ PARAMS
31
+ ------
32
+ C: compression factor used to compress
33
+ """
34
+ return torch.exp(x) / C
35
+
36
+
37
+ def spectral_normalize_torch(magnitudes):
38
+ output = dynamic_range_compression_torch(magnitudes)
39
+ return output
40
+
41
+
42
+ def spectral_de_normalize_torch(magnitudes):
43
+ output = dynamic_range_decompression_torch(magnitudes)
44
+ return output
45
+
46
+
47
+ mel_basis = {}
48
+ hann_window = {}
49
+
50
+
51
+ def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
52
+ if torch.min(y) < -1.0:
53
+ print("min value is ", torch.min(y))
54
+ if torch.max(y) > 1.0:
55
+ print("max value is ", torch.max(y))
56
+
57
+ global hann_window
58
+ dtype_device = str(y.dtype) + "_" + str(y.device)
59
+ wnsize_dtype_device = str(win_size) + "_" + dtype_device
60
+ if wnsize_dtype_device not in hann_window:
61
+ hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
62
+ dtype=y.dtype, device=y.device
63
+ )
64
+
65
+ y = torch.nn.functional.pad(
66
+ y.unsqueeze(1),
67
+ (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
68
+ mode="reflect",
69
+ )
70
+ y = y.squeeze(1)
71
+ spec = torch.stft(
72
+ y,
73
+ n_fft,
74
+ hop_length=hop_size,
75
+ win_length=win_size,
76
+ window=hann_window[wnsize_dtype_device],
77
+ center=center,
78
+ pad_mode="reflect",
79
+ normalized=False,
80
+ onesided=True,
81
+ return_complex=False,
82
+ )
83
+
84
+ spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
85
+ return spec
86
+
87
+
88
+ def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
89
+ global mel_basis
90
+ dtype_device = str(spec.dtype) + "_" + str(spec.device)
91
+ fmax_dtype_device = str(fmax) + "_" + dtype_device
92
+ if fmax_dtype_device not in mel_basis:
93
+ mel = librosa_mel_fn(
94
+ sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
95
+ )
96
+ mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
97
+ dtype=spec.dtype, device=spec.device
98
+ )
99
+ spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
100
+ spec = spectral_normalize_torch(spec)
101
+ return spec
102
+
103
+
104
+ def mel_spectrogram_torch(
105
+ y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
106
+ ):
107
+ if torch.min(y) < -1.0:
108
+ print("min value is ", torch.min(y))
109
+ if torch.max(y) > 1.0:
110
+ print("max value is ", torch.max(y))
111
+
112
+ global mel_basis, hann_window
113
+ dtype_device = str(y.dtype) + "_" + str(y.device)
114
+ fmax_dtype_device = str(fmax) + "_" + dtype_device
115
+ wnsize_dtype_device = str(win_size) + "_" + dtype_device
116
+ if fmax_dtype_device not in mel_basis:
117
+ mel = librosa_mel_fn(
118
+ sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
119
+ )
120
+ mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
121
+ dtype=y.dtype, device=y.device
122
+ )
123
+ if wnsize_dtype_device not in hann_window:
124
+ hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
125
+ dtype=y.dtype, device=y.device
126
+ )
127
+
128
+ y = torch.nn.functional.pad(
129
+ y.unsqueeze(1),
130
+ (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
131
+ mode="reflect",
132
+ )
133
+ y = y.squeeze(1)
134
+
135
+ spec = torch.stft(
136
+ y,
137
+ n_fft,
138
+ hop_length=hop_size,
139
+ win_length=win_size,
140
+ window=hann_window[wnsize_dtype_device],
141
+ center=center,
142
+ pad_mode="reflect",
143
+ normalized=False,
144
+ onesided=True,
145
+ return_complex=False,
146
+ )
147
+
148
+ spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
149
+
150
+ spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
151
+ spec = spectral_normalize_torch(spec)
152
+
153
+ return spec
lipsync-models/GPT_SoVITS/module/models.py ADDED
@@ -0,0 +1,993 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import math
3
+
4
+ import torch
5
+ from module import attentions, commons, modules
6
+ from module.commons import get_padding, init_weights
7
+ from module.mrte_model import MRTE
8
+ from module.quantize import ResidualVectorQuantizer
9
+ from text import symbols
10
+ from torch import nn
11
+ from torch.cuda.amp import autocast
12
+ from torch.nn import AvgPool1d, Conv1d, Conv2d, ConvTranspose1d
13
+ from torch.nn import functional as F
14
+ from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
15
+
16
+
17
+ class StochasticDurationPredictor(nn.Module):
18
+ def __init__(
19
+ self,
20
+ in_channels,
21
+ filter_channels,
22
+ kernel_size,
23
+ p_dropout,
24
+ n_flows=4,
25
+ gin_channels=0,
26
+ ):
27
+ super().__init__()
28
+ filter_channels = in_channels # it needs to be removed from future version.
29
+ self.in_channels = in_channels
30
+ self.filter_channels = filter_channels
31
+ self.kernel_size = kernel_size
32
+ self.p_dropout = p_dropout
33
+ self.n_flows = n_flows
34
+ self.gin_channels = gin_channels
35
+
36
+ self.log_flow = modules.Log()
37
+ self.flows = nn.ModuleList()
38
+ self.flows.append(modules.ElementwiseAffine(2))
39
+ for i in range(n_flows):
40
+ self.flows.append(
41
+ modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
42
+ )
43
+ self.flows.append(modules.Flip())
44
+
45
+ self.post_pre = nn.Conv1d(1, filter_channels, 1)
46
+ self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
47
+ self.post_convs = modules.DDSConv(
48
+ filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
49
+ )
50
+ self.post_flows = nn.ModuleList()
51
+ self.post_flows.append(modules.ElementwiseAffine(2))
52
+ for i in range(4):
53
+ self.post_flows.append(
54
+ modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3)
55
+ )
56
+ self.post_flows.append(modules.Flip())
57
+
58
+ self.pre = nn.Conv1d(in_channels, filter_channels, 1)
59
+ self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
60
+ self.convs = modules.DDSConv(
61
+ filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
62
+ )
63
+ if gin_channels != 0:
64
+ self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
65
+
66
+ def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
67
+ x = torch.detach(x)
68
+ x = self.pre(x)
69
+ if g is not None:
70
+ g = torch.detach(g)
71
+ x = x + self.cond(g)
72
+ x = self.convs(x, x_mask)
73
+ x = self.proj(x) * x_mask
74
+
75
+ if not reverse:
76
+ flows = self.flows
77
+ assert w is not None
78
+
79
+ logdet_tot_q = 0
80
+ h_w = self.post_pre(w)
81
+ h_w = self.post_convs(h_w, x_mask)
82
+ h_w = self.post_proj(h_w) * x_mask
83
+ e_q = (
84
+ torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
85
+ * x_mask
86
+ )
87
+ z_q = e_q
88
+ for flow in self.post_flows:
89
+ z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
90
+ logdet_tot_q += logdet_q
91
+ z_u, z1 = torch.split(z_q, [1, 1], 1)
92
+ u = torch.sigmoid(z_u) * x_mask
93
+ z0 = (w - u) * x_mask
94
+ logdet_tot_q += torch.sum(
95
+ (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
96
+ )
97
+ logq = (
98
+ torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
99
+ - logdet_tot_q
100
+ )
101
+
102
+ logdet_tot = 0
103
+ z0, logdet = self.log_flow(z0, x_mask)
104
+ logdet_tot += logdet
105
+ z = torch.cat([z0, z1], 1)
106
+ for flow in flows:
107
+ z, logdet = flow(z, x_mask, g=x, reverse=reverse)
108
+ logdet_tot = logdet_tot + logdet
109
+ nll = (
110
+ torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
111
+ - logdet_tot
112
+ )
113
+ return nll + logq # [b]
114
+ else:
115
+ flows = list(reversed(self.flows))
116
+ flows = flows[:-2] + [flows[-1]] # remove a useless vflow
117
+ z = (
118
+ torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
119
+ * noise_scale
120
+ )
121
+ for flow in flows:
122
+ z = flow(z, x_mask, g=x, reverse=reverse)
123
+ z0, z1 = torch.split(z, [1, 1], 1)
124
+ logw = z0
125
+ return logw
126
+
127
+
128
+ class DurationPredictor(nn.Module):
129
+ def __init__(
130
+ self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
131
+ ):
132
+ super().__init__()
133
+
134
+ self.in_channels = in_channels
135
+ self.filter_channels = filter_channels
136
+ self.kernel_size = kernel_size
137
+ self.p_dropout = p_dropout
138
+ self.gin_channels = gin_channels
139
+
140
+ self.drop = nn.Dropout(p_dropout)
141
+ self.conv_1 = nn.Conv1d(
142
+ in_channels, filter_channels, kernel_size, padding=kernel_size // 2
143
+ )
144
+ self.norm_1 = modules.LayerNorm(filter_channels)
145
+ self.conv_2 = nn.Conv1d(
146
+ filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
147
+ )
148
+ self.norm_2 = modules.LayerNorm(filter_channels)
149
+ self.proj = nn.Conv1d(filter_channels, 1, 1)
150
+
151
+ if gin_channels != 0:
152
+ self.cond = nn.Conv1d(gin_channels, in_channels, 1)
153
+
154
+ def forward(self, x, x_mask, g=None):
155
+ x = torch.detach(x)
156
+ if g is not None:
157
+ g = torch.detach(g)
158
+ x = x + self.cond(g)
159
+ x = self.conv_1(x * x_mask)
160
+ x = torch.relu(x)
161
+ x = self.norm_1(x)
162
+ x = self.drop(x)
163
+ x = self.conv_2(x * x_mask)
164
+ x = torch.relu(x)
165
+ x = self.norm_2(x)
166
+ x = self.drop(x)
167
+ x = self.proj(x * x_mask)
168
+ return x * x_mask
169
+
170
+
171
+ class TextEncoder(nn.Module):
172
+ def __init__(
173
+ self,
174
+ out_channels,
175
+ hidden_channels,
176
+ filter_channels,
177
+ n_heads,
178
+ n_layers,
179
+ kernel_size,
180
+ p_dropout,
181
+ latent_channels=192,
182
+ ):
183
+ super().__init__()
184
+ self.out_channels = out_channels
185
+ self.hidden_channels = hidden_channels
186
+ self.filter_channels = filter_channels
187
+ self.n_heads = n_heads
188
+ self.n_layers = n_layers
189
+ self.kernel_size = kernel_size
190
+ self.p_dropout = p_dropout
191
+ self.latent_channels = latent_channels
192
+
193
+ self.ssl_proj = nn.Conv1d(768, hidden_channels, 1)
194
+
195
+ self.encoder_ssl = attentions.Encoder(
196
+ hidden_channels,
197
+ filter_channels,
198
+ n_heads,
199
+ n_layers // 2,
200
+ kernel_size,
201
+ p_dropout,
202
+ )
203
+
204
+ self.encoder_text = attentions.Encoder(
205
+ hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
206
+ )
207
+ self.text_embedding = nn.Embedding(len(symbols), hidden_channels)
208
+
209
+ self.mrte = MRTE()
210
+
211
+ self.encoder2 = attentions.Encoder(
212
+ hidden_channels,
213
+ filter_channels,
214
+ n_heads,
215
+ n_layers // 2,
216
+ kernel_size,
217
+ p_dropout,
218
+ )
219
+
220
+ self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
221
+
222
+ def forward(self, y, y_lengths, text, text_lengths, ge, test=None):
223
+ y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(
224
+ y.dtype
225
+ )
226
+
227
+ y = self.ssl_proj(y * y_mask) * y_mask
228
+
229
+ y = self.encoder_ssl(y * y_mask, y_mask)
230
+
231
+ text_mask = torch.unsqueeze(
232
+ commons.sequence_mask(text_lengths, text.size(1)), 1
233
+ ).to(y.dtype)
234
+ if test == 1:
235
+ text[:, :] = 0
236
+ text = self.text_embedding(text).transpose(1, 2)
237
+ text = self.encoder_text(text * text_mask, text_mask)
238
+ y = self.mrte(y, y_mask, text, text_mask, ge)
239
+
240
+ y = self.encoder2(y * y_mask, y_mask)
241
+
242
+ stats = self.proj(y) * y_mask
243
+ m, logs = torch.split(stats, self.out_channels, dim=1)
244
+ return y, m, logs, y_mask
245
+
246
+ def extract_latent(self, x):
247
+ x = self.ssl_proj(x)
248
+ quantized, codes, commit_loss, quantized_list = self.quantizer(x)
249
+ return codes.transpose(0, 1)
250
+
251
+ def decode_latent(self, codes, y_mask, refer, refer_mask, ge):
252
+ quantized = self.quantizer.decode(codes)
253
+
254
+ y = self.vq_proj(quantized) * y_mask
255
+ y = self.encoder_ssl(y * y_mask, y_mask)
256
+
257
+ y = self.mrte(y, y_mask, refer, refer_mask, ge)
258
+
259
+ y = self.encoder2(y * y_mask, y_mask)
260
+
261
+ stats = self.proj(y) * y_mask
262
+ m, logs = torch.split(stats, self.out_channels, dim=1)
263
+ return y, m, logs, y_mask, quantized
264
+
265
+
266
+ class ResidualCouplingBlock(nn.Module):
267
+ def __init__(
268
+ self,
269
+ channels,
270
+ hidden_channels,
271
+ kernel_size,
272
+ dilation_rate,
273
+ n_layers,
274
+ n_flows=4,
275
+ gin_channels=0,
276
+ ):
277
+ super().__init__()
278
+ self.channels = channels
279
+ self.hidden_channels = hidden_channels
280
+ self.kernel_size = kernel_size
281
+ self.dilation_rate = dilation_rate
282
+ self.n_layers = n_layers
283
+ self.n_flows = n_flows
284
+ self.gin_channels = gin_channels
285
+
286
+ self.flows = nn.ModuleList()
287
+ for i in range(n_flows):
288
+ self.flows.append(
289
+ modules.ResidualCouplingLayer(
290
+ channels,
291
+ hidden_channels,
292
+ kernel_size,
293
+ dilation_rate,
294
+ n_layers,
295
+ gin_channels=gin_channels,
296
+ mean_only=True,
297
+ )
298
+ )
299
+ self.flows.append(modules.Flip())
300
+
301
+ def forward(self, x, x_mask, g=None, reverse=False):
302
+ if not reverse:
303
+ for flow in self.flows:
304
+ x, _ = flow(x, x_mask, g=g, reverse=reverse)
305
+ else:
306
+ for flow in reversed(self.flows):
307
+ x = flow(x, x_mask, g=g, reverse=reverse)
308
+ return x
309
+
310
+
311
+ class PosteriorEncoder(nn.Module):
312
+ def __init__(
313
+ self,
314
+ in_channels,
315
+ out_channels,
316
+ hidden_channels,
317
+ kernel_size,
318
+ dilation_rate,
319
+ n_layers,
320
+ gin_channels=0,
321
+ ):
322
+ super().__init__()
323
+ self.in_channels = in_channels
324
+ self.out_channels = out_channels
325
+ self.hidden_channels = hidden_channels
326
+ self.kernel_size = kernel_size
327
+ self.dilation_rate = dilation_rate
328
+ self.n_layers = n_layers
329
+ self.gin_channels = gin_channels
330
+
331
+ self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
332
+ self.enc = modules.WN(
333
+ hidden_channels,
334
+ kernel_size,
335
+ dilation_rate,
336
+ n_layers,
337
+ gin_channels=gin_channels,
338
+ )
339
+ self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
340
+
341
+ def forward(self, x, x_lengths, g=None):
342
+ if g != None:
343
+ g = g.detach()
344
+ x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
345
+ x.dtype
346
+ )
347
+ x = self.pre(x) * x_mask
348
+ x = self.enc(x, x_mask, g=g)
349
+ stats = self.proj(x) * x_mask
350
+ m, logs = torch.split(stats, self.out_channels, dim=1)
351
+ z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
352
+ return z, m, logs, x_mask
353
+
354
+
355
+ class WNEncoder(nn.Module):
356
+ def __init__(
357
+ self,
358
+ in_channels,
359
+ out_channels,
360
+ hidden_channels,
361
+ kernel_size,
362
+ dilation_rate,
363
+ n_layers,
364
+ gin_channels=0,
365
+ ):
366
+ super().__init__()
367
+ self.in_channels = in_channels
368
+ self.out_channels = out_channels
369
+ self.hidden_channels = hidden_channels
370
+ self.kernel_size = kernel_size
371
+ self.dilation_rate = dilation_rate
372
+ self.n_layers = n_layers
373
+ self.gin_channels = gin_channels
374
+
375
+ self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
376
+ self.enc = modules.WN(
377
+ hidden_channels,
378
+ kernel_size,
379
+ dilation_rate,
380
+ n_layers,
381
+ gin_channels=gin_channels,
382
+ )
383
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
384
+ self.norm = modules.LayerNorm(out_channels)
385
+
386
+ def forward(self, x, x_lengths, g=None):
387
+ x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
388
+ x.dtype
389
+ )
390
+ x = self.pre(x) * x_mask
391
+ x = self.enc(x, x_mask, g=g)
392
+ out = self.proj(x) * x_mask
393
+ out = self.norm(out)
394
+ return out
395
+
396
+
397
+ class Generator(torch.nn.Module):
398
+ def __init__(
399
+ self,
400
+ initial_channel,
401
+ resblock,
402
+ resblock_kernel_sizes,
403
+ resblock_dilation_sizes,
404
+ upsample_rates,
405
+ upsample_initial_channel,
406
+ upsample_kernel_sizes,
407
+ gin_channels=0,
408
+ ):
409
+ super(Generator, self).__init__()
410
+ self.num_kernels = len(resblock_kernel_sizes)
411
+ self.num_upsamples = len(upsample_rates)
412
+ self.conv_pre = Conv1d(
413
+ initial_channel, upsample_initial_channel, 7, 1, padding=3
414
+ )
415
+ resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
416
+
417
+ self.ups = nn.ModuleList()
418
+ for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
419
+ self.ups.append(
420
+ weight_norm(
421
+ ConvTranspose1d(
422
+ upsample_initial_channel // (2**i),
423
+ upsample_initial_channel // (2 ** (i + 1)),
424
+ k,
425
+ u,
426
+ padding=(k - u) // 2,
427
+ )
428
+ )
429
+ )
430
+
431
+ self.resblocks = nn.ModuleList()
432
+ for i in range(len(self.ups)):
433
+ ch = upsample_initial_channel // (2 ** (i + 1))
434
+ for j, (k, d) in enumerate(
435
+ zip(resblock_kernel_sizes, resblock_dilation_sizes)
436
+ ):
437
+ self.resblocks.append(resblock(ch, k, d))
438
+
439
+ self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
440
+ self.ups.apply(init_weights)
441
+
442
+ if gin_channels != 0:
443
+ self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
444
+
445
+ def forward(self, x, g=None):
446
+ x = self.conv_pre(x)
447
+ if g is not None:
448
+ x = x + self.cond(g)
449
+
450
+ for i in range(self.num_upsamples):
451
+ x = F.leaky_relu(x, modules.LRELU_SLOPE)
452
+ x = self.ups[i](x)
453
+ xs = None
454
+ for j in range(self.num_kernels):
455
+ if xs is None:
456
+ xs = self.resblocks[i * self.num_kernels + j](x)
457
+ else:
458
+ xs += self.resblocks[i * self.num_kernels + j](x)
459
+ x = xs / self.num_kernels
460
+ x = F.leaky_relu(x)
461
+ x = self.conv_post(x)
462
+ x = torch.tanh(x)
463
+
464
+ return x
465
+
466
+ def remove_weight_norm(self):
467
+ print("Removing weight norm...")
468
+ for l in self.ups:
469
+ remove_weight_norm(l)
470
+ for l in self.resblocks:
471
+ l.remove_weight_norm()
472
+
473
+
474
+ class DiscriminatorP(torch.nn.Module):
475
+ def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
476
+ super(DiscriminatorP, self).__init__()
477
+ self.period = period
478
+ self.use_spectral_norm = use_spectral_norm
479
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
480
+ self.convs = nn.ModuleList(
481
+ [
482
+ norm_f(
483
+ Conv2d(
484
+ 1,
485
+ 32,
486
+ (kernel_size, 1),
487
+ (stride, 1),
488
+ padding=(get_padding(kernel_size, 1), 0),
489
+ )
490
+ ),
491
+ norm_f(
492
+ Conv2d(
493
+ 32,
494
+ 128,
495
+ (kernel_size, 1),
496
+ (stride, 1),
497
+ padding=(get_padding(kernel_size, 1), 0),
498
+ )
499
+ ),
500
+ norm_f(
501
+ Conv2d(
502
+ 128,
503
+ 512,
504
+ (kernel_size, 1),
505
+ (stride, 1),
506
+ padding=(get_padding(kernel_size, 1), 0),
507
+ )
508
+ ),
509
+ norm_f(
510
+ Conv2d(
511
+ 512,
512
+ 1024,
513
+ (kernel_size, 1),
514
+ (stride, 1),
515
+ padding=(get_padding(kernel_size, 1), 0),
516
+ )
517
+ ),
518
+ norm_f(
519
+ Conv2d(
520
+ 1024,
521
+ 1024,
522
+ (kernel_size, 1),
523
+ 1,
524
+ padding=(get_padding(kernel_size, 1), 0),
525
+ )
526
+ ),
527
+ ]
528
+ )
529
+ self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
530
+
531
+ def forward(self, x):
532
+ fmap = []
533
+
534
+ # 1d to 2d
535
+ b, c, t = x.shape
536
+ if t % self.period != 0: # pad first
537
+ n_pad = self.period - (t % self.period)
538
+ x = F.pad(x, (0, n_pad), "reflect")
539
+ t = t + n_pad
540
+ x = x.view(b, c, t // self.period, self.period)
541
+
542
+ for l in self.convs:
543
+ x = l(x)
544
+ x = F.leaky_relu(x, modules.LRELU_SLOPE)
545
+ fmap.append(x)
546
+ x = self.conv_post(x)
547
+ fmap.append(x)
548
+ x = torch.flatten(x, 1, -1)
549
+
550
+ return x, fmap
551
+
552
+
553
+ class DiscriminatorS(torch.nn.Module):
554
+ def __init__(self, use_spectral_norm=False):
555
+ super(DiscriminatorS, self).__init__()
556
+ norm_f = weight_norm if use_spectral_norm == False else spectral_norm
557
+ self.convs = nn.ModuleList(
558
+ [
559
+ norm_f(Conv1d(1, 16, 15, 1, padding=7)),
560
+ norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
561
+ norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
562
+ norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
563
+ norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
564
+ norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
565
+ ]
566
+ )
567
+ self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
568
+
569
+ def forward(self, x):
570
+ fmap = []
571
+
572
+ for l in self.convs:
573
+ x = l(x)
574
+ x = F.leaky_relu(x, modules.LRELU_SLOPE)
575
+ fmap.append(x)
576
+ x = self.conv_post(x)
577
+ fmap.append(x)
578
+ x = torch.flatten(x, 1, -1)
579
+
580
+ return x, fmap
581
+
582
+
583
+ class MultiPeriodDiscriminator(torch.nn.Module):
584
+ def __init__(self, use_spectral_norm=False):
585
+ super(MultiPeriodDiscriminator, self).__init__()
586
+ periods = [2, 3, 5, 7, 11]
587
+
588
+ discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
589
+ discs = discs + [
590
+ DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods
591
+ ]
592
+ self.discriminators = nn.ModuleList(discs)
593
+
594
+ def forward(self, y, y_hat):
595
+ y_d_rs = []
596
+ y_d_gs = []
597
+ fmap_rs = []
598
+ fmap_gs = []
599
+ for i, d in enumerate(self.discriminators):
600
+ y_d_r, fmap_r = d(y)
601
+ y_d_g, fmap_g = d(y_hat)
602
+ y_d_rs.append(y_d_r)
603
+ y_d_gs.append(y_d_g)
604
+ fmap_rs.append(fmap_r)
605
+ fmap_gs.append(fmap_g)
606
+
607
+ return y_d_rs, y_d_gs, fmap_rs, fmap_gs
608
+
609
+
610
+ class ReferenceEncoder(nn.Module):
611
+ """
612
+ inputs --- [N, Ty/r, n_mels*r] mels
613
+ outputs --- [N, ref_enc_gru_size]
614
+ """
615
+
616
+ def __init__(self, spec_channels, gin_channels=0):
617
+ super().__init__()
618
+ self.spec_channels = spec_channels
619
+ ref_enc_filters = [32, 32, 64, 64, 128, 128]
620
+ K = len(ref_enc_filters)
621
+ filters = [1] + ref_enc_filters
622
+ convs = [
623
+ weight_norm(
624
+ nn.Conv2d(
625
+ in_channels=filters[i],
626
+ out_channels=filters[i + 1],
627
+ kernel_size=(3, 3),
628
+ stride=(2, 2),
629
+ padding=(1, 1),
630
+ )
631
+ )
632
+ for i in range(K)
633
+ ]
634
+ self.convs = nn.ModuleList(convs)
635
+ # self.wns = nn.ModuleList([weight_norm(num_features=ref_enc_filters[i]) for i in range(K)])
636
+
637
+ out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
638
+ self.gru = nn.GRU(
639
+ input_size=ref_enc_filters[-1] * out_channels,
640
+ hidden_size=256 // 2,
641
+ batch_first=True,
642
+ )
643
+ self.proj = nn.Linear(128, gin_channels)
644
+
645
+ def forward(self, inputs):
646
+ N = inputs.size(0)
647
+ out = inputs.view(N, 1, -1, self.spec_channels) # [N, 1, Ty, n_freqs]
648
+ for conv in self.convs:
649
+ out = conv(out)
650
+ # out = wn(out)
651
+ out = F.relu(out) # [N, 128, Ty//2^K, n_mels//2^K]
652
+
653
+ out = out.transpose(1, 2) # [N, Ty//2^K, 128, n_mels//2^K]
654
+ T = out.size(1)
655
+ N = out.size(0)
656
+ out = out.contiguous().view(N, T, -1) # [N, Ty//2^K, 128*n_mels//2^K]
657
+
658
+ self.gru.flatten_parameters()
659
+ memory, out = self.gru(out) # out --- [1, N, 128]
660
+
661
+ return self.proj(out.squeeze(0)).unsqueeze(-1)
662
+
663
+ def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
664
+ for i in range(n_convs):
665
+ L = (L - kernel_size + 2 * pad) // stride + 1
666
+ return L
667
+
668
+
669
+ class Quantizer_module(torch.nn.Module):
670
+ def __init__(self, n_e, e_dim):
671
+ super(Quantizer_module, self).__init__()
672
+ self.embedding = nn.Embedding(n_e, e_dim)
673
+ self.embedding.weight.data.uniform_(-1.0 / n_e, 1.0 / n_e)
674
+
675
+ def forward(self, x):
676
+ d = (
677
+ torch.sum(x**2, 1, keepdim=True)
678
+ + torch.sum(self.embedding.weight**2, 1)
679
+ - 2 * torch.matmul(x, self.embedding.weight.T)
680
+ )
681
+ min_indicies = torch.argmin(d, 1)
682
+ z_q = self.embedding(min_indicies)
683
+ return z_q, min_indicies
684
+
685
+
686
+ class Quantizer(torch.nn.Module):
687
+ def __init__(self, embed_dim=512, n_code_groups=4, n_codes=160):
688
+ super(Quantizer, self).__init__()
689
+ assert embed_dim % n_code_groups == 0
690
+ self.quantizer_modules = nn.ModuleList(
691
+ [
692
+ Quantizer_module(n_codes, embed_dim // n_code_groups)
693
+ for _ in range(n_code_groups)
694
+ ]
695
+ )
696
+ self.n_code_groups = n_code_groups
697
+ self.embed_dim = embed_dim
698
+
699
+ def forward(self, xin):
700
+ # B, C, T
701
+ B, C, T = xin.shape
702
+ xin = xin.transpose(1, 2)
703
+ x = xin.reshape(-1, self.embed_dim)
704
+ x = torch.split(x, self.embed_dim // self.n_code_groups, dim=-1)
705
+ min_indicies = []
706
+ z_q = []
707
+ for _x, m in zip(x, self.quantizer_modules):
708
+ _z_q, _min_indicies = m(_x)
709
+ z_q.append(_z_q)
710
+ min_indicies.append(_min_indicies) # B * T,
711
+ z_q = torch.cat(z_q, -1).reshape(xin.shape)
712
+ loss = 0.25 * torch.mean((z_q.detach() - xin) ** 2) + torch.mean(
713
+ (z_q - xin.detach()) ** 2
714
+ )
715
+ z_q = xin + (z_q - xin).detach()
716
+ z_q = z_q.transpose(1, 2)
717
+ codes = torch.stack(min_indicies, -1).reshape(B, T, self.n_code_groups)
718
+ return z_q, loss, codes.transpose(1, 2)
719
+
720
+ def embed(self, x):
721
+ # idx: N, 4, T
722
+ x = x.transpose(1, 2)
723
+ x = torch.split(x, 1, 2)
724
+ ret = []
725
+ for q, embed in zip(x, self.quantizer_modules):
726
+ q = embed.embedding(q.squeeze(-1))
727
+ ret.append(q)
728
+ ret = torch.cat(ret, -1)
729
+ return ret.transpose(1, 2) # N, C, T
730
+
731
+
732
+ class CodePredictor(nn.Module):
733
+ def __init__(
734
+ self,
735
+ hidden_channels,
736
+ filter_channels,
737
+ n_heads,
738
+ n_layers,
739
+ kernel_size,
740
+ p_dropout,
741
+ n_q=8,
742
+ dims=1024,
743
+ ssl_dim=768,
744
+ ):
745
+ super().__init__()
746
+ self.hidden_channels = hidden_channels
747
+ self.filter_channels = filter_channels
748
+ self.n_heads = n_heads
749
+ self.n_layers = n_layers
750
+ self.kernel_size = kernel_size
751
+ self.p_dropout = p_dropout
752
+
753
+ self.vq_proj = nn.Conv1d(ssl_dim, hidden_channels, 1)
754
+ self.ref_enc = modules.MelStyleEncoder(
755
+ ssl_dim, style_vector_dim=hidden_channels
756
+ )
757
+
758
+ self.encoder = attentions.Encoder(
759
+ hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
760
+ )
761
+
762
+ self.out_proj = nn.Conv1d(hidden_channels, (n_q - 1) * dims, 1)
763
+ self.n_q = n_q
764
+ self.dims = dims
765
+
766
+ def forward(self, x, x_mask, refer, codes, infer=False):
767
+ x = x.detach()
768
+ x = self.vq_proj(x * x_mask) * x_mask
769
+ g = self.ref_enc(refer, x_mask)
770
+ x = x + g
771
+ x = self.encoder(x * x_mask, x_mask)
772
+ x = self.out_proj(x * x_mask) * x_mask
773
+ logits = x.reshape(x.shape[0], self.n_q - 1, self.dims, x.shape[-1]).transpose(
774
+ 2, 3
775
+ )
776
+ target = codes[1:].transpose(0, 1)
777
+ if not infer:
778
+ logits = logits.reshape(-1, self.dims)
779
+ target = target.reshape(-1)
780
+ loss = torch.nn.functional.cross_entropy(logits, target)
781
+ return loss
782
+ else:
783
+ _, top10_preds = torch.topk(logits, 10, dim=-1)
784
+ correct_top10 = torch.any(top10_preds == target.unsqueeze(-1), dim=-1)
785
+ top3_acc = 100 * torch.mean(correct_top10.float()).detach().cpu().item()
786
+
787
+ print("Top-10 Accuracy:", top3_acc, "%")
788
+
789
+ pred_codes = torch.argmax(logits, dim=-1)
790
+ acc = 100 * torch.mean((pred_codes == target).float()).detach().cpu().item()
791
+ print("Top-1 Accuracy:", acc, "%")
792
+
793
+ return pred_codes.transpose(0, 1)
794
+
795
+
796
+ class SynthesizerTrn(nn.Module):
797
+ """
798
+ Synthesizer for Training
799
+ """
800
+
801
+ def __init__(
802
+ self,
803
+ spec_channels,
804
+ segment_size,
805
+ inter_channels,
806
+ hidden_channels,
807
+ filter_channels,
808
+ n_heads,
809
+ n_layers,
810
+ kernel_size,
811
+ p_dropout,
812
+ resblock,
813
+ resblock_kernel_sizes,
814
+ resblock_dilation_sizes,
815
+ upsample_rates,
816
+ upsample_initial_channel,
817
+ upsample_kernel_sizes,
818
+ n_speakers=0,
819
+ gin_channels=0,
820
+ use_sdp=True,
821
+ semantic_frame_rate=None,
822
+ freeze_quantizer=None,
823
+ **kwargs
824
+ ):
825
+ super().__init__()
826
+ self.spec_channels = spec_channels
827
+ self.inter_channels = inter_channels
828
+ self.hidden_channels = hidden_channels
829
+ self.filter_channels = filter_channels
830
+ self.n_heads = n_heads
831
+ self.n_layers = n_layers
832
+ self.kernel_size = kernel_size
833
+ self.p_dropout = p_dropout
834
+ self.resblock = resblock
835
+ self.resblock_kernel_sizes = resblock_kernel_sizes
836
+ self.resblock_dilation_sizes = resblock_dilation_sizes
837
+ self.upsample_rates = upsample_rates
838
+ self.upsample_initial_channel = upsample_initial_channel
839
+ self.upsample_kernel_sizes = upsample_kernel_sizes
840
+ self.segment_size = segment_size
841
+ self.n_speakers = n_speakers
842
+ self.gin_channels = gin_channels
843
+
844
+ self.use_sdp = use_sdp
845
+ self.enc_p = TextEncoder(
846
+ inter_channels,
847
+ hidden_channels,
848
+ filter_channels,
849
+ n_heads,
850
+ n_layers,
851
+ kernel_size,
852
+ p_dropout,
853
+ )
854
+ self.dec = Generator(
855
+ inter_channels,
856
+ resblock,
857
+ resblock_kernel_sizes,
858
+ resblock_dilation_sizes,
859
+ upsample_rates,
860
+ upsample_initial_channel,
861
+ upsample_kernel_sizes,
862
+ gin_channels=gin_channels,
863
+ )
864
+ self.enc_q = PosteriorEncoder(
865
+ spec_channels,
866
+ inter_channels,
867
+ hidden_channels,
868
+ 5,
869
+ 1,
870
+ 16,
871
+ gin_channels=gin_channels,
872
+ )
873
+ self.flow = ResidualCouplingBlock(
874
+ inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels
875
+ )
876
+
877
+ self.ref_enc = modules.MelStyleEncoder(
878
+ spec_channels, style_vector_dim=gin_channels
879
+ )
880
+
881
+ ssl_dim = 768
882
+ assert semantic_frame_rate in ["25hz", "50hz"]
883
+ self.semantic_frame_rate = semantic_frame_rate
884
+ if semantic_frame_rate == "25hz":
885
+ self.ssl_proj = nn.Conv1d(ssl_dim, ssl_dim, 2, stride=2)
886
+ else:
887
+ self.ssl_proj = nn.Conv1d(ssl_dim, ssl_dim, 1, stride=1)
888
+
889
+ self.quantizer = ResidualVectorQuantizer(dimension=ssl_dim, n_q=1, bins=1024)
890
+ if freeze_quantizer:
891
+ self.ssl_proj.requires_grad_(False)
892
+ self.quantizer.requires_grad_(False)
893
+ # self.enc_p.text_embedding.requires_grad_(False)
894
+ # self.enc_p.encoder_text.requires_grad_(False)
895
+ # self.enc_p.mrte.requires_grad_(False)
896
+
897
+ def forward(self, ssl, y, y_lengths, text, text_lengths):
898
+ y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(
899
+ y.dtype
900
+ )
901
+ ge = self.ref_enc(y * y_mask, y_mask)
902
+
903
+ with autocast(enabled=False):
904
+ ssl = self.ssl_proj(ssl)
905
+ quantized, codes, commit_loss, quantized_list = self.quantizer(
906
+ ssl, layers=[0]
907
+ )
908
+
909
+ if self.semantic_frame_rate == "25hz":
910
+ quantized = F.interpolate(
911
+ quantized, size=int(quantized.shape[-1] * 2), mode="nearest"
912
+ )
913
+
914
+ x, m_p, logs_p, y_mask = self.enc_p(
915
+ quantized, y_lengths, text, text_lengths, ge
916
+ )
917
+ z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=ge)
918
+ z_p = self.flow(z, y_mask, g=ge)
919
+
920
+ z_slice, ids_slice = commons.rand_slice_segments(
921
+ z, y_lengths, self.segment_size
922
+ )
923
+ o = self.dec(z_slice, g=ge)
924
+ return (
925
+ o,
926
+ commit_loss,
927
+ ids_slice,
928
+ y_mask,
929
+ y_mask,
930
+ (z, z_p, m_p, logs_p, m_q, logs_q),
931
+ quantized,
932
+ )
933
+
934
+ def infer(self, ssl, y, y_lengths, text, text_lengths, test=None, noise_scale=0.5):
935
+ y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(
936
+ y.dtype
937
+ )
938
+ ge = self.ref_enc(y * y_mask, y_mask)
939
+
940
+ ssl = self.ssl_proj(ssl)
941
+ quantized, codes, commit_loss, _ = self.quantizer(ssl, layers=[0])
942
+ if self.semantic_frame_rate == "25hz":
943
+ quantized = F.interpolate(
944
+ quantized, size=int(quantized.shape[-1] * 2), mode="nearest"
945
+ )
946
+
947
+ x, m_p, logs_p, y_mask = self.enc_p(
948
+ quantized, y_lengths, text, text_lengths, ge, test=test
949
+ )
950
+ z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
951
+
952
+ z = self.flow(z_p, y_mask, g=ge, reverse=True)
953
+
954
+ o = self.dec((z * y_mask)[:, :, :], g=ge)
955
+ return o, y_mask, (z, z_p, m_p, logs_p)
956
+
957
+ @torch.no_grad()
958
+ def decode_by_ge(self, codes, text, ge=None, noise_scale=0.5):
959
+ y_lengths = torch.LongTensor([codes.size(2) * 2]).to(codes.device)
960
+ text_lengths = torch.LongTensor([text.size(-1)]).to(text.device)
961
+
962
+ quantized = self.quantizer.decode(codes)
963
+ if self.semantic_frame_rate == "25hz":
964
+ quantized = F.interpolate(
965
+ quantized, size=int(quantized.shape[-1] * 2), mode="nearest"
966
+ )
967
+
968
+ x, m_p, logs_p, y_mask = self.enc_p(
969
+ quantized, y_lengths, text, text_lengths, ge
970
+ )
971
+ z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
972
+
973
+ z = self.flow(z_p, y_mask, g=ge, reverse=True)
974
+
975
+ o = self.dec((z * y_mask)[:, :, :], g=ge)
976
+ return o
977
+
978
+ @torch.no_grad()
979
+ def decode(self, codes, text, refer, noise_scale=0.5):
980
+ ge = None
981
+ if refer is not None:
982
+ refer_lengths = torch.LongTensor([refer.size(2)]).to(refer.device)
983
+ refer_mask = torch.unsqueeze(
984
+ commons.sequence_mask(refer_lengths, refer.size(2)), 1
985
+ ).to(refer.dtype)
986
+ ge = self.ref_enc(refer * refer_mask, refer_mask)
987
+
988
+ return self.decode_by_ge(codes, text, ge, noise_scale)
989
+
990
+ def extract_latent(self, x):
991
+ ssl = self.ssl_proj(x)
992
+ quantized, codes, commit_loss, quantized_list = self.quantizer(ssl)
993
+ return codes.transpose(0, 1)
lipsync-models/GPT_SoVITS/module/modules.py ADDED
@@ -0,0 +1,923 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import numpy as np
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn import functional as F
6
+
7
+ from torch.nn import Conv1d
8
+ from torch.nn.utils import weight_norm, remove_weight_norm
9
+
10
+ from module import commons
11
+ from module.commons import init_weights, get_padding
12
+ from module.transforms import piecewise_rational_quadratic_transform
13
+ import torch.distributions as D
14
+
15
+
16
+ LRELU_SLOPE = 0.1
17
+
18
+
19
+ class LayerNorm(nn.Module):
20
+ def __init__(self, channels, eps=1e-5):
21
+ super().__init__()
22
+ self.channels = channels
23
+ self.eps = eps
24
+
25
+ self.gamma = nn.Parameter(torch.ones(channels))
26
+ self.beta = nn.Parameter(torch.zeros(channels))
27
+
28
+ def forward(self, x):
29
+ x = x.transpose(1, -1)
30
+ x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
31
+ return x.transpose(1, -1)
32
+
33
+
34
+ class ConvReluNorm(nn.Module):
35
+ def __init__(
36
+ self,
37
+ in_channels,
38
+ hidden_channels,
39
+ out_channels,
40
+ kernel_size,
41
+ n_layers,
42
+ p_dropout,
43
+ ):
44
+ super().__init__()
45
+ self.in_channels = in_channels
46
+ self.hidden_channels = hidden_channels
47
+ self.out_channels = out_channels
48
+ self.kernel_size = kernel_size
49
+ self.n_layers = n_layers
50
+ self.p_dropout = p_dropout
51
+ assert n_layers > 1, "Number of layers should be larger than 0."
52
+
53
+ self.conv_layers = nn.ModuleList()
54
+ self.norm_layers = nn.ModuleList()
55
+ self.conv_layers.append(
56
+ nn.Conv1d(
57
+ in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
58
+ )
59
+ )
60
+ self.norm_layers.append(LayerNorm(hidden_channels))
61
+ self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
62
+ for _ in range(n_layers - 1):
63
+ self.conv_layers.append(
64
+ nn.Conv1d(
65
+ hidden_channels,
66
+ hidden_channels,
67
+ kernel_size,
68
+ padding=kernel_size // 2,
69
+ )
70
+ )
71
+ self.norm_layers.append(LayerNorm(hidden_channels))
72
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
73
+ self.proj.weight.data.zero_()
74
+ self.proj.bias.data.zero_()
75
+
76
+ def forward(self, x, x_mask):
77
+ x_org = x
78
+ for i in range(self.n_layers):
79
+ x = self.conv_layers[i](x * x_mask)
80
+ x = self.norm_layers[i](x)
81
+ x = self.relu_drop(x)
82
+ x = x_org + self.proj(x)
83
+ return x * x_mask
84
+
85
+
86
+ class DDSConv(nn.Module):
87
+ """
88
+ Dialted and Depth-Separable Convolution
89
+ """
90
+
91
+ def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
92
+ super().__init__()
93
+ self.channels = channels
94
+ self.kernel_size = kernel_size
95
+ self.n_layers = n_layers
96
+ self.p_dropout = p_dropout
97
+
98
+ self.drop = nn.Dropout(p_dropout)
99
+ self.convs_sep = nn.ModuleList()
100
+ self.convs_1x1 = nn.ModuleList()
101
+ self.norms_1 = nn.ModuleList()
102
+ self.norms_2 = nn.ModuleList()
103
+ for i in range(n_layers):
104
+ dilation = kernel_size**i
105
+ padding = (kernel_size * dilation - dilation) // 2
106
+ self.convs_sep.append(
107
+ nn.Conv1d(
108
+ channels,
109
+ channels,
110
+ kernel_size,
111
+ groups=channels,
112
+ dilation=dilation,
113
+ padding=padding,
114
+ )
115
+ )
116
+ self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
117
+ self.norms_1.append(LayerNorm(channels))
118
+ self.norms_2.append(LayerNorm(channels))
119
+
120
+ def forward(self, x, x_mask, g=None):
121
+ if g is not None:
122
+ x = x + g
123
+ for i in range(self.n_layers):
124
+ y = self.convs_sep[i](x * x_mask)
125
+ y = self.norms_1[i](y)
126
+ y = F.gelu(y)
127
+ y = self.convs_1x1[i](y)
128
+ y = self.norms_2[i](y)
129
+ y = F.gelu(y)
130
+ y = self.drop(y)
131
+ x = x + y
132
+ return x * x_mask
133
+
134
+
135
+ class WN(torch.nn.Module):
136
+ def __init__(
137
+ self,
138
+ hidden_channels,
139
+ kernel_size,
140
+ dilation_rate,
141
+ n_layers,
142
+ gin_channels=0,
143
+ p_dropout=0,
144
+ ):
145
+ super(WN, self).__init__()
146
+ assert kernel_size % 2 == 1
147
+ self.hidden_channels = hidden_channels
148
+ self.kernel_size = (kernel_size,)
149
+ self.dilation_rate = dilation_rate
150
+ self.n_layers = n_layers
151
+ self.gin_channels = gin_channels
152
+ self.p_dropout = p_dropout
153
+
154
+ self.in_layers = torch.nn.ModuleList()
155
+ self.res_skip_layers = torch.nn.ModuleList()
156
+ self.drop = nn.Dropout(p_dropout)
157
+
158
+ if gin_channels != 0:
159
+ cond_layer = torch.nn.Conv1d(
160
+ gin_channels, 2 * hidden_channels * n_layers, 1
161
+ )
162
+ self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
163
+
164
+ for i in range(n_layers):
165
+ dilation = dilation_rate**i
166
+ padding = int((kernel_size * dilation - dilation) / 2)
167
+ in_layer = torch.nn.Conv1d(
168
+ hidden_channels,
169
+ 2 * hidden_channels,
170
+ kernel_size,
171
+ dilation=dilation,
172
+ padding=padding,
173
+ )
174
+ in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
175
+ self.in_layers.append(in_layer)
176
+
177
+ # last one is not necessary
178
+ if i < n_layers - 1:
179
+ res_skip_channels = 2 * hidden_channels
180
+ else:
181
+ res_skip_channels = hidden_channels
182
+
183
+ res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
184
+ res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
185
+ self.res_skip_layers.append(res_skip_layer)
186
+
187
+ def forward(self, x, x_mask, g=None, **kwargs):
188
+ output = torch.zeros_like(x)
189
+ n_channels_tensor = torch.IntTensor([self.hidden_channels])
190
+
191
+ if g is not None:
192
+ g = self.cond_layer(g)
193
+
194
+ for i in range(self.n_layers):
195
+ x_in = self.in_layers[i](x)
196
+ if g is not None:
197
+ cond_offset = i * 2 * self.hidden_channels
198
+ g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
199
+ else:
200
+ g_l = torch.zeros_like(x_in)
201
+
202
+ acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
203
+ acts = self.drop(acts)
204
+
205
+ res_skip_acts = self.res_skip_layers[i](acts)
206
+ if i < self.n_layers - 1:
207
+ res_acts = res_skip_acts[:, : self.hidden_channels, :]
208
+ x = (x + res_acts) * x_mask
209
+ output = output + res_skip_acts[:, self.hidden_channels :, :]
210
+ else:
211
+ output = output + res_skip_acts
212
+ return output * x_mask
213
+
214
+ def remove_weight_norm(self):
215
+ if self.gin_channels != 0:
216
+ torch.nn.utils.remove_weight_norm(self.cond_layer)
217
+ for l in self.in_layers:
218
+ torch.nn.utils.remove_weight_norm(l)
219
+ for l in self.res_skip_layers:
220
+ torch.nn.utils.remove_weight_norm(l)
221
+
222
+
223
+ class ResBlock1(torch.nn.Module):
224
+ def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
225
+ super(ResBlock1, self).__init__()
226
+ self.convs1 = nn.ModuleList(
227
+ [
228
+ weight_norm(
229
+ Conv1d(
230
+ channels,
231
+ channels,
232
+ kernel_size,
233
+ 1,
234
+ dilation=dilation[0],
235
+ padding=get_padding(kernel_size, dilation[0]),
236
+ )
237
+ ),
238
+ weight_norm(
239
+ Conv1d(
240
+ channels,
241
+ channels,
242
+ kernel_size,
243
+ 1,
244
+ dilation=dilation[1],
245
+ padding=get_padding(kernel_size, dilation[1]),
246
+ )
247
+ ),
248
+ weight_norm(
249
+ Conv1d(
250
+ channels,
251
+ channels,
252
+ kernel_size,
253
+ 1,
254
+ dilation=dilation[2],
255
+ padding=get_padding(kernel_size, dilation[2]),
256
+ )
257
+ ),
258
+ ]
259
+ )
260
+ self.convs1.apply(init_weights)
261
+
262
+ self.convs2 = nn.ModuleList(
263
+ [
264
+ weight_norm(
265
+ Conv1d(
266
+ channels,
267
+ channels,
268
+ kernel_size,
269
+ 1,
270
+ dilation=1,
271
+ padding=get_padding(kernel_size, 1),
272
+ )
273
+ ),
274
+ weight_norm(
275
+ Conv1d(
276
+ channels,
277
+ channels,
278
+ kernel_size,
279
+ 1,
280
+ dilation=1,
281
+ padding=get_padding(kernel_size, 1),
282
+ )
283
+ ),
284
+ weight_norm(
285
+ Conv1d(
286
+ channels,
287
+ channels,
288
+ kernel_size,
289
+ 1,
290
+ dilation=1,
291
+ padding=get_padding(kernel_size, 1),
292
+ )
293
+ ),
294
+ ]
295
+ )
296
+ self.convs2.apply(init_weights)
297
+
298
+ def forward(self, x, x_mask=None):
299
+ for c1, c2 in zip(self.convs1, self.convs2):
300
+ xt = F.leaky_relu(x, LRELU_SLOPE)
301
+ if x_mask is not None:
302
+ xt = xt * x_mask
303
+ xt = c1(xt)
304
+ xt = F.leaky_relu(xt, LRELU_SLOPE)
305
+ if x_mask is not None:
306
+ xt = xt * x_mask
307
+ xt = c2(xt)
308
+ x = xt + x
309
+ if x_mask is not None:
310
+ x = x * x_mask
311
+ return x
312
+
313
+ def remove_weight_norm(self):
314
+ for l in self.convs1:
315
+ remove_weight_norm(l)
316
+ for l in self.convs2:
317
+ remove_weight_norm(l)
318
+
319
+
320
+ class ResBlock2(torch.nn.Module):
321
+ def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
322
+ super(ResBlock2, self).__init__()
323
+ self.convs = nn.ModuleList(
324
+ [
325
+ weight_norm(
326
+ Conv1d(
327
+ channels,
328
+ channels,
329
+ kernel_size,
330
+ 1,
331
+ dilation=dilation[0],
332
+ padding=get_padding(kernel_size, dilation[0]),
333
+ )
334
+ ),
335
+ weight_norm(
336
+ Conv1d(
337
+ channels,
338
+ channels,
339
+ kernel_size,
340
+ 1,
341
+ dilation=dilation[1],
342
+ padding=get_padding(kernel_size, dilation[1]),
343
+ )
344
+ ),
345
+ ]
346
+ )
347
+ self.convs.apply(init_weights)
348
+
349
+ def forward(self, x, x_mask=None):
350
+ for c in self.convs:
351
+ xt = F.leaky_relu(x, LRELU_SLOPE)
352
+ if x_mask is not None:
353
+ xt = xt * x_mask
354
+ xt = c(xt)
355
+ x = xt + x
356
+ if x_mask is not None:
357
+ x = x * x_mask
358
+ return x
359
+
360
+ def remove_weight_norm(self):
361
+ for l in self.convs:
362
+ remove_weight_norm(l)
363
+
364
+
365
+ class Log(nn.Module):
366
+ def forward(self, x, x_mask, reverse=False, **kwargs):
367
+ if not reverse:
368
+ y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
369
+ logdet = torch.sum(-y, [1, 2])
370
+ return y, logdet
371
+ else:
372
+ x = torch.exp(x) * x_mask
373
+ return x
374
+
375
+
376
+ class Flip(nn.Module):
377
+ def forward(self, x, *args, reverse=False, **kwargs):
378
+ x = torch.flip(x, [1])
379
+ if not reverse:
380
+ logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
381
+ return x, logdet
382
+ else:
383
+ return x
384
+
385
+
386
+ class ElementwiseAffine(nn.Module):
387
+ def __init__(self, channels):
388
+ super().__init__()
389
+ self.channels = channels
390
+ self.m = nn.Parameter(torch.zeros(channels, 1))
391
+ self.logs = nn.Parameter(torch.zeros(channels, 1))
392
+
393
+ def forward(self, x, x_mask, reverse=False, **kwargs):
394
+ if not reverse:
395
+ y = self.m + torch.exp(self.logs) * x
396
+ y = y * x_mask
397
+ logdet = torch.sum(self.logs * x_mask, [1, 2])
398
+ return y, logdet
399
+ else:
400
+ x = (x - self.m) * torch.exp(-self.logs) * x_mask
401
+ return x
402
+
403
+
404
+ class ResidualCouplingLayer(nn.Module):
405
+ def __init__(
406
+ self,
407
+ channels,
408
+ hidden_channels,
409
+ kernel_size,
410
+ dilation_rate,
411
+ n_layers,
412
+ p_dropout=0,
413
+ gin_channels=0,
414
+ mean_only=False,
415
+ ):
416
+ assert channels % 2 == 0, "channels should be divisible by 2"
417
+ super().__init__()
418
+ self.channels = channels
419
+ self.hidden_channels = hidden_channels
420
+ self.kernel_size = kernel_size
421
+ self.dilation_rate = dilation_rate
422
+ self.n_layers = n_layers
423
+ self.half_channels = channels // 2
424
+ self.mean_only = mean_only
425
+
426
+ self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
427
+ self.enc = WN(
428
+ hidden_channels,
429
+ kernel_size,
430
+ dilation_rate,
431
+ n_layers,
432
+ p_dropout=p_dropout,
433
+ gin_channels=gin_channels,
434
+ )
435
+ self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
436
+ self.post.weight.data.zero_()
437
+ self.post.bias.data.zero_()
438
+
439
+ def forward(self, x, x_mask, g=None, reverse=False):
440
+ x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
441
+ h = self.pre(x0) * x_mask
442
+ h = self.enc(h, x_mask, g=g)
443
+ stats = self.post(h) * x_mask
444
+ if not self.mean_only:
445
+ m, logs = torch.split(stats, [self.half_channels] * 2, 1)
446
+ else:
447
+ m = stats
448
+ logs = torch.zeros_like(m)
449
+
450
+ if not reverse:
451
+ x1 = m + x1 * torch.exp(logs) * x_mask
452
+ x = torch.cat([x0, x1], 1)
453
+ logdet = torch.sum(logs, [1, 2])
454
+ return x, logdet
455
+ else:
456
+ x1 = (x1 - m) * torch.exp(-logs) * x_mask
457
+ x = torch.cat([x0, x1], 1)
458
+ return x
459
+
460
+
461
+ class ConvFlow(nn.Module):
462
+ def __init__(
463
+ self,
464
+ in_channels,
465
+ filter_channels,
466
+ kernel_size,
467
+ n_layers,
468
+ num_bins=10,
469
+ tail_bound=5.0,
470
+ ):
471
+ super().__init__()
472
+ self.in_channels = in_channels
473
+ self.filter_channels = filter_channels
474
+ self.kernel_size = kernel_size
475
+ self.n_layers = n_layers
476
+ self.num_bins = num_bins
477
+ self.tail_bound = tail_bound
478
+ self.half_channels = in_channels // 2
479
+
480
+ self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
481
+ self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
482
+ self.proj = nn.Conv1d(
483
+ filter_channels, self.half_channels * (num_bins * 3 - 1), 1
484
+ )
485
+ self.proj.weight.data.zero_()
486
+ self.proj.bias.data.zero_()
487
+
488
+ def forward(self, x, x_mask, g=None, reverse=False):
489
+ x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
490
+ h = self.pre(x0)
491
+ h = self.convs(h, x_mask, g=g)
492
+ h = self.proj(h) * x_mask
493
+
494
+ b, c, t = x0.shape
495
+ h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
496
+
497
+ unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
498
+ unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
499
+ self.filter_channels
500
+ )
501
+ unnormalized_derivatives = h[..., 2 * self.num_bins :]
502
+
503
+ x1, logabsdet = piecewise_rational_quadratic_transform(
504
+ x1,
505
+ unnormalized_widths,
506
+ unnormalized_heights,
507
+ unnormalized_derivatives,
508
+ inverse=reverse,
509
+ tails="linear",
510
+ tail_bound=self.tail_bound,
511
+ )
512
+
513
+ x = torch.cat([x0, x1], 1) * x_mask
514
+ logdet = torch.sum(logabsdet * x_mask, [1, 2])
515
+ if not reverse:
516
+ return x, logdet
517
+ else:
518
+ return x
519
+
520
+
521
+ class LinearNorm(nn.Module):
522
+ def __init__(
523
+ self,
524
+ in_channels,
525
+ out_channels,
526
+ bias=True,
527
+ spectral_norm=False,
528
+ ):
529
+ super(LinearNorm, self).__init__()
530
+ self.fc = nn.Linear(in_channels, out_channels, bias)
531
+
532
+ if spectral_norm:
533
+ self.fc = nn.utils.spectral_norm(self.fc)
534
+
535
+ def forward(self, input):
536
+ out = self.fc(input)
537
+ return out
538
+
539
+
540
+ class Mish(nn.Module):
541
+ def __init__(self):
542
+ super(Mish, self).__init__()
543
+
544
+ def forward(self, x):
545
+ return x * torch.tanh(F.softplus(x))
546
+
547
+
548
+ class Conv1dGLU(nn.Module):
549
+ """
550
+ Conv1d + GLU(Gated Linear Unit) with residual connection.
551
+ For GLU refer to https://arxiv.org/abs/1612.08083 paper.
552
+ """
553
+
554
+ def __init__(self, in_channels, out_channels, kernel_size, dropout):
555
+ super(Conv1dGLU, self).__init__()
556
+ self.out_channels = out_channels
557
+ self.conv1 = ConvNorm(in_channels, 2 * out_channels, kernel_size=kernel_size)
558
+ self.dropout = nn.Dropout(dropout)
559
+
560
+ def forward(self, x):
561
+ residual = x
562
+ x = self.conv1(x)
563
+ x1, x2 = torch.split(x, split_size_or_sections=self.out_channels, dim=1)
564
+ x = x1 * torch.sigmoid(x2)
565
+ x = residual + self.dropout(x)
566
+ return x
567
+
568
+
569
+ class ConvNorm(nn.Module):
570
+ def __init__(
571
+ self,
572
+ in_channels,
573
+ out_channels,
574
+ kernel_size=1,
575
+ stride=1,
576
+ padding=None,
577
+ dilation=1,
578
+ bias=True,
579
+ spectral_norm=False,
580
+ ):
581
+ super(ConvNorm, self).__init__()
582
+
583
+ if padding is None:
584
+ assert kernel_size % 2 == 1
585
+ padding = int(dilation * (kernel_size - 1) / 2)
586
+
587
+ self.conv = torch.nn.Conv1d(
588
+ in_channels,
589
+ out_channels,
590
+ kernel_size=kernel_size,
591
+ stride=stride,
592
+ padding=padding,
593
+ dilation=dilation,
594
+ bias=bias,
595
+ )
596
+
597
+ if spectral_norm:
598
+ self.conv = nn.utils.spectral_norm(self.conv)
599
+
600
+ def forward(self, input):
601
+ out = self.conv(input)
602
+ return out
603
+
604
+
605
+ class MultiHeadAttention(nn.Module):
606
+ """Multi-Head Attention module"""
607
+
608
+ def __init__(self, n_head, d_model, d_k, d_v, dropout=0.0, spectral_norm=False):
609
+ super().__init__()
610
+
611
+ self.n_head = n_head
612
+ self.d_k = d_k
613
+ self.d_v = d_v
614
+
615
+ self.w_qs = nn.Linear(d_model, n_head * d_k)
616
+ self.w_ks = nn.Linear(d_model, n_head * d_k)
617
+ self.w_vs = nn.Linear(d_model, n_head * d_v)
618
+
619
+ self.attention = ScaledDotProductAttention(
620
+ temperature=np.power(d_model, 0.5), dropout=dropout
621
+ )
622
+
623
+ self.fc = nn.Linear(n_head * d_v, d_model)
624
+ self.dropout = nn.Dropout(dropout)
625
+
626
+ if spectral_norm:
627
+ self.w_qs = nn.utils.spectral_norm(self.w_qs)
628
+ self.w_ks = nn.utils.spectral_norm(self.w_ks)
629
+ self.w_vs = nn.utils.spectral_norm(self.w_vs)
630
+ self.fc = nn.utils.spectral_norm(self.fc)
631
+
632
+ def forward(self, x, mask=None):
633
+ d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
634
+ sz_b, len_x, _ = x.size()
635
+
636
+ residual = x
637
+
638
+ q = self.w_qs(x).view(sz_b, len_x, n_head, d_k)
639
+ k = self.w_ks(x).view(sz_b, len_x, n_head, d_k)
640
+ v = self.w_vs(x).view(sz_b, len_x, n_head, d_v)
641
+ q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_k) # (n*b) x lq x dk
642
+ k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_k) # (n*b) x lk x dk
643
+ v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_v) # (n*b) x lv x dv
644
+
645
+ if mask is not None:
646
+ slf_mask = mask.repeat(n_head, 1, 1) # (n*b) x .. x ..
647
+ else:
648
+ slf_mask = None
649
+ output, attn = self.attention(q, k, v, mask=slf_mask)
650
+
651
+ output = output.view(n_head, sz_b, len_x, d_v)
652
+ output = (
653
+ output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_x, -1)
654
+ ) # b x lq x (n*dv)
655
+
656
+ output = self.fc(output)
657
+
658
+ output = self.dropout(output) + residual
659
+ return output, attn
660
+
661
+
662
+ class ScaledDotProductAttention(nn.Module):
663
+ """Scaled Dot-Product Attention"""
664
+
665
+ def __init__(self, temperature, dropout):
666
+ super().__init__()
667
+ self.temperature = temperature
668
+ self.softmax = nn.Softmax(dim=2)
669
+ self.dropout = nn.Dropout(dropout)
670
+
671
+ def forward(self, q, k, v, mask=None):
672
+ attn = torch.bmm(q, k.transpose(1, 2))
673
+ attn = attn / self.temperature
674
+
675
+ if mask is not None:
676
+ attn = attn.masked_fill(mask, -np.inf)
677
+
678
+ attn = self.softmax(attn)
679
+ p_attn = self.dropout(attn)
680
+
681
+ output = torch.bmm(p_attn, v)
682
+ return output, attn
683
+
684
+
685
+ class MelStyleEncoder(nn.Module):
686
+ """MelStyleEncoder"""
687
+
688
+ def __init__(
689
+ self,
690
+ n_mel_channels=80,
691
+ style_hidden=128,
692
+ style_vector_dim=256,
693
+ style_kernel_size=5,
694
+ style_head=2,
695
+ dropout=0.1,
696
+ ):
697
+ super(MelStyleEncoder, self).__init__()
698
+ self.in_dim = n_mel_channels
699
+ self.hidden_dim = style_hidden
700
+ self.out_dim = style_vector_dim
701
+ self.kernel_size = style_kernel_size
702
+ self.n_head = style_head
703
+ self.dropout = dropout
704
+
705
+ self.spectral = nn.Sequential(
706
+ LinearNorm(self.in_dim, self.hidden_dim),
707
+ Mish(),
708
+ nn.Dropout(self.dropout),
709
+ LinearNorm(self.hidden_dim, self.hidden_dim),
710
+ Mish(),
711
+ nn.Dropout(self.dropout),
712
+ )
713
+
714
+ self.temporal = nn.Sequential(
715
+ Conv1dGLU(self.hidden_dim, self.hidden_dim, self.kernel_size, self.dropout),
716
+ Conv1dGLU(self.hidden_dim, self.hidden_dim, self.kernel_size, self.dropout),
717
+ )
718
+
719
+ self.slf_attn = MultiHeadAttention(
720
+ self.n_head,
721
+ self.hidden_dim,
722
+ self.hidden_dim // self.n_head,
723
+ self.hidden_dim // self.n_head,
724
+ self.dropout,
725
+ )
726
+
727
+ self.fc = LinearNorm(self.hidden_dim, self.out_dim)
728
+
729
+ def temporal_avg_pool(self, x, mask=None):
730
+ if mask is None:
731
+ out = torch.mean(x, dim=1)
732
+ else:
733
+ len_ = (~mask).sum(dim=1).unsqueeze(1)
734
+ x = x.masked_fill(mask.unsqueeze(-1), 0)
735
+ x = x.sum(dim=1)
736
+ out = torch.div(x, len_)
737
+ return out
738
+
739
+ def forward(self, x, mask=None):
740
+ x = x.transpose(1, 2)
741
+ if mask is not None:
742
+ mask = (mask.int() == 0).squeeze(1)
743
+ max_len = x.shape[1]
744
+ slf_attn_mask = (
745
+ mask.unsqueeze(1).expand(-1, max_len, -1) if mask is not None else None
746
+ )
747
+
748
+ # spectral
749
+ x = self.spectral(x)
750
+ # temporal
751
+ x = x.transpose(1, 2)
752
+ x = self.temporal(x)
753
+ x = x.transpose(1, 2)
754
+ # self-attention
755
+ if mask is not None:
756
+ x = x.masked_fill(mask.unsqueeze(-1), 0)
757
+ x, _ = self.slf_attn(x, mask=slf_attn_mask)
758
+ # fc
759
+ x = self.fc(x)
760
+ # temoral average pooling
761
+ w = self.temporal_avg_pool(x, mask=mask)
762
+
763
+ return w.unsqueeze(-1)
764
+
765
+
766
+ class MelStyleEncoderVAE(nn.Module):
767
+ def __init__(self, spec_channels, z_latent_dim, emb_dim):
768
+ super().__init__()
769
+ self.ref_encoder = MelStyleEncoder(spec_channels, style_vector_dim=emb_dim)
770
+ self.fc1 = nn.Linear(emb_dim, z_latent_dim)
771
+ self.fc2 = nn.Linear(emb_dim, z_latent_dim)
772
+ self.fc3 = nn.Linear(z_latent_dim, emb_dim)
773
+ self.z_latent_dim = z_latent_dim
774
+
775
+ def reparameterize(self, mu, logvar):
776
+ if self.training:
777
+ std = torch.exp(0.5 * logvar)
778
+ eps = torch.randn_like(std)
779
+ return eps.mul(std).add_(mu)
780
+ else:
781
+ return mu
782
+
783
+ def forward(self, inputs, mask=None):
784
+ enc_out = self.ref_encoder(inputs.squeeze(-1), mask).squeeze(-1)
785
+ mu = self.fc1(enc_out)
786
+ logvar = self.fc2(enc_out)
787
+ posterior = D.Normal(mu, torch.exp(logvar))
788
+ kl_divergence = D.kl_divergence(
789
+ posterior, D.Normal(torch.zeros_like(mu), torch.ones_like(logvar))
790
+ )
791
+ loss_kl = kl_divergence.mean()
792
+
793
+ z = posterior.rsample()
794
+ style_embed = self.fc3(z)
795
+
796
+ return style_embed.unsqueeze(-1), loss_kl
797
+
798
+ def infer(self, inputs=None, random_sample=False, manual_latent=None):
799
+ if manual_latent is None:
800
+ if random_sample:
801
+ dev = next(self.parameters()).device
802
+ posterior = D.Normal(
803
+ torch.zeros(1, self.z_latent_dim, device=dev),
804
+ torch.ones(1, self.z_latent_dim, device=dev),
805
+ )
806
+ z = posterior.rsample()
807
+ else:
808
+ enc_out = self.ref_encoder(inputs.transpose(1, 2))
809
+ mu = self.fc1(enc_out)
810
+ z = mu
811
+ else:
812
+ z = manual_latent
813
+ style_embed = self.fc3(z)
814
+ return style_embed.unsqueeze(-1), z
815
+
816
+
817
+ class ActNorm(nn.Module):
818
+ def __init__(self, channels, ddi=False, **kwargs):
819
+ super().__init__()
820
+ self.channels = channels
821
+ self.initialized = not ddi
822
+
823
+ self.logs = nn.Parameter(torch.zeros(1, channels, 1))
824
+ self.bias = nn.Parameter(torch.zeros(1, channels, 1))
825
+
826
+ def forward(self, x, x_mask=None, g=None, reverse=False, **kwargs):
827
+ if x_mask is None:
828
+ x_mask = torch.ones(x.size(0), 1, x.size(2)).to(
829
+ device=x.device, dtype=x.dtype
830
+ )
831
+ x_len = torch.sum(x_mask, [1, 2])
832
+ if not self.initialized:
833
+ self.initialize(x, x_mask)
834
+ self.initialized = True
835
+
836
+ if reverse:
837
+ z = (x - self.bias) * torch.exp(-self.logs) * x_mask
838
+ logdet = None
839
+ return z
840
+ else:
841
+ z = (self.bias + torch.exp(self.logs) * x) * x_mask
842
+ logdet = torch.sum(self.logs) * x_len # [b]
843
+ return z, logdet
844
+
845
+ def store_inverse(self):
846
+ pass
847
+
848
+ def set_ddi(self, ddi):
849
+ self.initialized = not ddi
850
+
851
+ def initialize(self, x, x_mask):
852
+ with torch.no_grad():
853
+ denom = torch.sum(x_mask, [0, 2])
854
+ m = torch.sum(x * x_mask, [0, 2]) / denom
855
+ m_sq = torch.sum(x * x * x_mask, [0, 2]) / denom
856
+ v = m_sq - (m**2)
857
+ logs = 0.5 * torch.log(torch.clamp_min(v, 1e-6))
858
+
859
+ bias_init = (
860
+ (-m * torch.exp(-logs)).view(*self.bias.shape).to(dtype=self.bias.dtype)
861
+ )
862
+ logs_init = (-logs).view(*self.logs.shape).to(dtype=self.logs.dtype)
863
+
864
+ self.bias.data.copy_(bias_init)
865
+ self.logs.data.copy_(logs_init)
866
+
867
+
868
+ class InvConvNear(nn.Module):
869
+ def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
870
+ super().__init__()
871
+ assert n_split % 2 == 0
872
+ self.channels = channels
873
+ self.n_split = n_split
874
+ self.no_jacobian = no_jacobian
875
+
876
+ w_init = torch.linalg.qr(
877
+ torch.FloatTensor(self.n_split, self.n_split).normal_()
878
+ )[0]
879
+ if torch.det(w_init) < 0:
880
+ w_init[:, 0] = -1 * w_init[:, 0]
881
+ self.weight = nn.Parameter(w_init)
882
+
883
+ def forward(self, x, x_mask=None, g=None, reverse=False, **kwargs):
884
+ b, c, t = x.size()
885
+ assert c % self.n_split == 0
886
+ if x_mask is None:
887
+ x_mask = 1
888
+ x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t
889
+ else:
890
+ x_len = torch.sum(x_mask, [1, 2])
891
+
892
+ x = x.view(b, 2, c // self.n_split, self.n_split // 2, t)
893
+ x = (
894
+ x.permute(0, 1, 3, 2, 4)
895
+ .contiguous()
896
+ .view(b, self.n_split, c // self.n_split, t)
897
+ )
898
+
899
+ if reverse:
900
+ if hasattr(self, "weight_inv"):
901
+ weight = self.weight_inv
902
+ else:
903
+ weight = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
904
+ logdet = None
905
+ else:
906
+ weight = self.weight
907
+ if self.no_jacobian:
908
+ logdet = 0
909
+ else:
910
+ logdet = torch.logdet(self.weight) * (c / self.n_split) * x_len # [b]
911
+
912
+ weight = weight.view(self.n_split, self.n_split, 1, 1)
913
+ z = F.conv2d(x, weight)
914
+
915
+ z = z.view(b, 2, self.n_split // 2, c // self.n_split, t)
916
+ z = z.permute(0, 1, 3, 2, 4).contiguous().view(b, c, t) * x_mask
917
+ if reverse:
918
+ return z
919
+ else:
920
+ return z, logdet
921
+
922
+ def store_inverse(self):
923
+ self.weight_inv = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
lipsync-models/GPT_SoVITS/module/mrte_model.py ADDED
@@ -0,0 +1,192 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This is Multi-reference timbre encoder
2
+
3
+ import torch
4
+ from torch import nn
5
+ from torch.nn.utils import remove_weight_norm, weight_norm
6
+ from module.attentions import MultiHeadAttention
7
+
8
+
9
+ class MRTE(nn.Module):
10
+ def __init__(
11
+ self,
12
+ content_enc_channels=192,
13
+ hidden_size=512,
14
+ out_channels=192,
15
+ kernel_size=5,
16
+ n_heads=4,
17
+ ge_layer=2,
18
+ ):
19
+ super(MRTE, self).__init__()
20
+ self.cross_attention = MultiHeadAttention(hidden_size, hidden_size, n_heads)
21
+ self.c_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)
22
+ self.text_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)
23
+ self.c_post = nn.Conv1d(hidden_size, out_channels, 1)
24
+
25
+ def forward(self, ssl_enc, ssl_mask, text, text_mask, ge, test=None):
26
+ if ge == None:
27
+ ge = 0
28
+ attn_mask = text_mask.unsqueeze(2) * ssl_mask.unsqueeze(-1)
29
+
30
+ ssl_enc = self.c_pre(ssl_enc * ssl_mask)
31
+ text_enc = self.text_pre(text * text_mask)
32
+ if test != None:
33
+ if test == 0:
34
+ x = (
35
+ self.cross_attention(
36
+ ssl_enc * ssl_mask, text_enc * text_mask, attn_mask
37
+ )
38
+ + ssl_enc
39
+ + ge
40
+ )
41
+ elif test == 1:
42
+ x = ssl_enc + ge
43
+ elif test == 2:
44
+ x = (
45
+ self.cross_attention(
46
+ ssl_enc * 0 * ssl_mask, text_enc * text_mask, attn_mask
47
+ )
48
+ + ge
49
+ )
50
+ else:
51
+ raise ValueError("test should be 0,1,2")
52
+ else:
53
+ x = (
54
+ self.cross_attention(
55
+ ssl_enc * ssl_mask, text_enc * text_mask, attn_mask
56
+ )
57
+ + ssl_enc
58
+ + ge
59
+ )
60
+ x = self.c_post(x * ssl_mask)
61
+ return x
62
+
63
+
64
+ class SpeakerEncoder(torch.nn.Module):
65
+ def __init__(
66
+ self,
67
+ mel_n_channels=80,
68
+ model_num_layers=2,
69
+ model_hidden_size=256,
70
+ model_embedding_size=256,
71
+ ):
72
+ super(SpeakerEncoder, self).__init__()
73
+ self.lstm = nn.LSTM(
74
+ mel_n_channels, model_hidden_size, model_num_layers, batch_first=True
75
+ )
76
+ self.linear = nn.Linear(model_hidden_size, model_embedding_size)
77
+ self.relu = nn.ReLU()
78
+
79
+ def forward(self, mels):
80
+ self.lstm.flatten_parameters()
81
+ _, (hidden, _) = self.lstm(mels.transpose(-1, -2))
82
+ embeds_raw = self.relu(self.linear(hidden[-1]))
83
+ return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
84
+
85
+
86
+ class MELEncoder(nn.Module):
87
+ def __init__(
88
+ self,
89
+ in_channels,
90
+ out_channels,
91
+ hidden_channels,
92
+ kernel_size,
93
+ dilation_rate,
94
+ n_layers,
95
+ ):
96
+ super().__init__()
97
+ self.in_channels = in_channels
98
+ self.out_channels = out_channels
99
+ self.hidden_channels = hidden_channels
100
+ self.kernel_size = kernel_size
101
+ self.dilation_rate = dilation_rate
102
+ self.n_layers = n_layers
103
+
104
+ self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
105
+ self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers)
106
+ self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
107
+
108
+ def forward(self, x):
109
+ # print(x.shape,x_lengths.shape)
110
+ x = self.pre(x)
111
+ x = self.enc(x)
112
+ x = self.proj(x)
113
+ return x
114
+
115
+
116
+ class WN(torch.nn.Module):
117
+ def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers):
118
+ super(WN, self).__init__()
119
+ assert kernel_size % 2 == 1
120
+ self.hidden_channels = hidden_channels
121
+ self.kernel_size = kernel_size
122
+ self.dilation_rate = dilation_rate
123
+ self.n_layers = n_layers
124
+
125
+ self.in_layers = torch.nn.ModuleList()
126
+ self.res_skip_layers = torch.nn.ModuleList()
127
+
128
+ for i in range(n_layers):
129
+ dilation = dilation_rate**i
130
+ padding = int((kernel_size * dilation - dilation) / 2)
131
+ in_layer = nn.Conv1d(
132
+ hidden_channels,
133
+ 2 * hidden_channels,
134
+ kernel_size,
135
+ dilation=dilation,
136
+ padding=padding,
137
+ )
138
+ in_layer = weight_norm(in_layer)
139
+ self.in_layers.append(in_layer)
140
+
141
+ # last one is not necessary
142
+ if i < n_layers - 1:
143
+ res_skip_channels = 2 * hidden_channels
144
+ else:
145
+ res_skip_channels = hidden_channels
146
+
147
+ res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
148
+ res_skip_layer = weight_norm(res_skip_layer, name="weight")
149
+ self.res_skip_layers.append(res_skip_layer)
150
+
151
+ def forward(self, x):
152
+ output = torch.zeros_like(x)
153
+ n_channels_tensor = torch.IntTensor([self.hidden_channels])
154
+
155
+ for i in range(self.n_layers):
156
+ x_in = self.in_layers[i](x)
157
+
158
+ acts = fused_add_tanh_sigmoid_multiply(x_in, n_channels_tensor)
159
+
160
+ res_skip_acts = self.res_skip_layers[i](acts)
161
+ if i < self.n_layers - 1:
162
+ res_acts = res_skip_acts[:, : self.hidden_channels, :]
163
+ x = x + res_acts
164
+ output = output + res_skip_acts[:, self.hidden_channels :, :]
165
+ else:
166
+ output = output + res_skip_acts
167
+ return output
168
+
169
+ def remove_weight_norm(self):
170
+ for l in self.in_layers:
171
+ remove_weight_norm(l)
172
+ for l in self.res_skip_layers:
173
+ remove_weight_norm(l)
174
+
175
+
176
+ @torch.jit.script
177
+ def fused_add_tanh_sigmoid_multiply(input, n_channels):
178
+ n_channels_int = n_channels[0]
179
+ t_act = torch.tanh(input[:, :n_channels_int, :])
180
+ s_act = torch.sigmoid(input[:, n_channels_int:, :])
181
+ acts = t_act * s_act
182
+ return acts
183
+
184
+
185
+ if __name__ == "__main__":
186
+ content_enc = torch.randn(3, 192, 100)
187
+ content_mask = torch.ones(3, 1, 100)
188
+ ref_mel = torch.randn(3, 128, 30)
189
+ ref_mask = torch.ones(3, 1, 30)
190
+ model = MRTE()
191
+ out = model(content_enc, content_mask, ref_mel, ref_mask)
192
+ print(out.shape)
lipsync-models/GPT_SoVITS/module/quantize.py ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Residual vector quantizer implementation."""
8
+
9
+ from dataclasses import dataclass, field
10
+ import math
11
+ import typing as tp
12
+
13
+ import torch
14
+ from torch import nn
15
+
16
+ from module.core_vq import ResidualVectorQuantization
17
+
18
+
19
+ @dataclass
20
+ class QuantizedResult:
21
+ quantized: torch.Tensor
22
+ codes: torch.Tensor
23
+ bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item.
24
+ penalty: tp.Optional[torch.Tensor] = None
25
+ metrics: dict = field(default_factory=dict)
26
+
27
+
28
+ class ResidualVectorQuantizer(nn.Module):
29
+ """Residual Vector Quantizer.
30
+ Args:
31
+ dimension (int): Dimension of the codebooks.
32
+ n_q (int): Number of residual vector quantizers used.
33
+ bins (int): Codebook size.
34
+ decay (float): Decay for exponential moving average over the codebooks.
35
+ kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
36
+ kmeans_iters (int): Number of iterations used for kmeans initialization.
37
+ threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
38
+ that have an exponential moving average cluster size less than the specified threshold with
39
+ randomly selected vector from the current batch.
40
+ """
41
+
42
+ def __init__(
43
+ self,
44
+ dimension: int = 256,
45
+ n_q: int = 8,
46
+ bins: int = 1024,
47
+ decay: float = 0.99,
48
+ kmeans_init: bool = True,
49
+ kmeans_iters: int = 50,
50
+ threshold_ema_dead_code: int = 2,
51
+ ):
52
+ super().__init__()
53
+ self.n_q = n_q
54
+ self.dimension = dimension
55
+ self.bins = bins
56
+ self.decay = decay
57
+ self.kmeans_init = kmeans_init
58
+ self.kmeans_iters = kmeans_iters
59
+ self.threshold_ema_dead_code = threshold_ema_dead_code
60
+ self.vq = ResidualVectorQuantization(
61
+ dim=self.dimension,
62
+ codebook_size=self.bins,
63
+ num_quantizers=self.n_q,
64
+ decay=self.decay,
65
+ kmeans_init=self.kmeans_init,
66
+ kmeans_iters=self.kmeans_iters,
67
+ threshold_ema_dead_code=self.threshold_ema_dead_code,
68
+ )
69
+
70
+ def forward(
71
+ self,
72
+ x: torch.Tensor,
73
+ n_q: tp.Optional[int] = None,
74
+ layers: tp.Optional[list] = None,
75
+ ) -> QuantizedResult:
76
+ """Residual vector quantization on the given input tensor.
77
+ Args:
78
+ x (torch.Tensor): Input tensor.
79
+ n_q (int): Number of quantizer used to quantize. Default: All quantizers.
80
+ layers (list): Layer that need to return quantized. Defalt: None.
81
+ Returns:
82
+ QuantizedResult:
83
+ The quantized (or approximately quantized) representation with
84
+ the associated numbert quantizers and layer quantized required to return.
85
+ """
86
+ n_q = n_q if n_q else self.n_q
87
+ if layers and max(layers) >= n_q:
88
+ raise ValueError(
89
+ f"Last layer index in layers: A {max(layers)}. Number of quantizers in RVQ: B {self.n_q}. A must less than B."
90
+ )
91
+ quantized, codes, commit_loss, quantized_list = self.vq(
92
+ x, n_q=n_q, layers=layers
93
+ )
94
+ return quantized, codes, torch.mean(commit_loss), quantized_list
95
+
96
+ def encode(
97
+ self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None
98
+ ) -> torch.Tensor:
99
+ """Encode a given input tensor with the specified sample rate at the given bandwidth.
100
+ The RVQ encode method sets the appropriate number of quantizer to use
101
+ and returns indices for each quantizer.
102
+ Args:
103
+ x (torch.Tensor): Input tensor.
104
+ n_q (int): Number of quantizer used to quantize. Default: All quantizers.
105
+ st (int): Start to encode input from which layers. Default: 0.
106
+ """
107
+ n_q = n_q if n_q else self.n_q
108
+ st = st or 0
109
+ codes = self.vq.encode(x, n_q=n_q, st=st)
110
+ return codes
111
+
112
+ def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:
113
+ """Decode the given codes to the quantized representation.
114
+ Args:
115
+ codes (torch.Tensor): Input indices for each quantizer.
116
+ st (int): Start to decode input codes from which layers. Default: 0.
117
+ """
118
+ quantized = self.vq.decode(codes, st=st)
119
+ return quantized
lipsync-models/GPT_SoVITS/module/transforms.py ADDED
@@ -0,0 +1,209 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch.nn import functional as F
3
+
4
+ import numpy as np
5
+
6
+
7
+ DEFAULT_MIN_BIN_WIDTH = 1e-3
8
+ DEFAULT_MIN_BIN_HEIGHT = 1e-3
9
+ DEFAULT_MIN_DERIVATIVE = 1e-3
10
+
11
+
12
+ def piecewise_rational_quadratic_transform(
13
+ inputs,
14
+ unnormalized_widths,
15
+ unnormalized_heights,
16
+ unnormalized_derivatives,
17
+ inverse=False,
18
+ tails=None,
19
+ tail_bound=1.0,
20
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
21
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
22
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
23
+ ):
24
+ if tails is None:
25
+ spline_fn = rational_quadratic_spline
26
+ spline_kwargs = {}
27
+ else:
28
+ spline_fn = unconstrained_rational_quadratic_spline
29
+ spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
30
+
31
+ outputs, logabsdet = spline_fn(
32
+ inputs=inputs,
33
+ unnormalized_widths=unnormalized_widths,
34
+ unnormalized_heights=unnormalized_heights,
35
+ unnormalized_derivatives=unnormalized_derivatives,
36
+ inverse=inverse,
37
+ min_bin_width=min_bin_width,
38
+ min_bin_height=min_bin_height,
39
+ min_derivative=min_derivative,
40
+ **spline_kwargs
41
+ )
42
+ return outputs, logabsdet
43
+
44
+
45
+ def searchsorted(bin_locations, inputs, eps=1e-6):
46
+ bin_locations[..., -1] += eps
47
+ return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
48
+
49
+
50
+ def unconstrained_rational_quadratic_spline(
51
+ inputs,
52
+ unnormalized_widths,
53
+ unnormalized_heights,
54
+ unnormalized_derivatives,
55
+ inverse=False,
56
+ tails="linear",
57
+ tail_bound=1.0,
58
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
59
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
60
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
61
+ ):
62
+ inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
63
+ outside_interval_mask = ~inside_interval_mask
64
+
65
+ outputs = torch.zeros_like(inputs)
66
+ logabsdet = torch.zeros_like(inputs)
67
+
68
+ if tails == "linear":
69
+ unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
70
+ constant = np.log(np.exp(1 - min_derivative) - 1)
71
+ unnormalized_derivatives[..., 0] = constant
72
+ unnormalized_derivatives[..., -1] = constant
73
+
74
+ outputs[outside_interval_mask] = inputs[outside_interval_mask]
75
+ logabsdet[outside_interval_mask] = 0
76
+ else:
77
+ raise RuntimeError("{} tails are not implemented.".format(tails))
78
+
79
+ (
80
+ outputs[inside_interval_mask],
81
+ logabsdet[inside_interval_mask],
82
+ ) = rational_quadratic_spline(
83
+ inputs=inputs[inside_interval_mask],
84
+ unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
85
+ unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
86
+ unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
87
+ inverse=inverse,
88
+ left=-tail_bound,
89
+ right=tail_bound,
90
+ bottom=-tail_bound,
91
+ top=tail_bound,
92
+ min_bin_width=min_bin_width,
93
+ min_bin_height=min_bin_height,
94
+ min_derivative=min_derivative,
95
+ )
96
+
97
+ return outputs, logabsdet
98
+
99
+
100
+ def rational_quadratic_spline(
101
+ inputs,
102
+ unnormalized_widths,
103
+ unnormalized_heights,
104
+ unnormalized_derivatives,
105
+ inverse=False,
106
+ left=0.0,
107
+ right=1.0,
108
+ bottom=0.0,
109
+ top=1.0,
110
+ min_bin_width=DEFAULT_MIN_BIN_WIDTH,
111
+ min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
112
+ min_derivative=DEFAULT_MIN_DERIVATIVE,
113
+ ):
114
+ if torch.min(inputs) < left or torch.max(inputs) > right:
115
+ raise ValueError("Input to a transform is not within its domain")
116
+
117
+ num_bins = unnormalized_widths.shape[-1]
118
+
119
+ if min_bin_width * num_bins > 1.0:
120
+ raise ValueError("Minimal bin width too large for the number of bins")
121
+ if min_bin_height * num_bins > 1.0:
122
+ raise ValueError("Minimal bin height too large for the number of bins")
123
+
124
+ widths = F.softmax(unnormalized_widths, dim=-1)
125
+ widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
126
+ cumwidths = torch.cumsum(widths, dim=-1)
127
+ cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
128
+ cumwidths = (right - left) * cumwidths + left
129
+ cumwidths[..., 0] = left
130
+ cumwidths[..., -1] = right
131
+ widths = cumwidths[..., 1:] - cumwidths[..., :-1]
132
+
133
+ derivatives = min_derivative + F.softplus(unnormalized_derivatives)
134
+
135
+ heights = F.softmax(unnormalized_heights, dim=-1)
136
+ heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
137
+ cumheights = torch.cumsum(heights, dim=-1)
138
+ cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
139
+ cumheights = (top - bottom) * cumheights + bottom
140
+ cumheights[..., 0] = bottom
141
+ cumheights[..., -1] = top
142
+ heights = cumheights[..., 1:] - cumheights[..., :-1]
143
+
144
+ if inverse:
145
+ bin_idx = searchsorted(cumheights, inputs)[..., None]
146
+ else:
147
+ bin_idx = searchsorted(cumwidths, inputs)[..., None]
148
+
149
+ input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
150
+ input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
151
+
152
+ input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
153
+ delta = heights / widths
154
+ input_delta = delta.gather(-1, bin_idx)[..., 0]
155
+
156
+ input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
157
+ input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
158
+
159
+ input_heights = heights.gather(-1, bin_idx)[..., 0]
160
+
161
+ if inverse:
162
+ a = (inputs - input_cumheights) * (
163
+ input_derivatives + input_derivatives_plus_one - 2 * input_delta
164
+ ) + input_heights * (input_delta - input_derivatives)
165
+ b = input_heights * input_derivatives - (inputs - input_cumheights) * (
166
+ input_derivatives + input_derivatives_plus_one - 2 * input_delta
167
+ )
168
+ c = -input_delta * (inputs - input_cumheights)
169
+
170
+ discriminant = b.pow(2) - 4 * a * c
171
+ assert (discriminant >= 0).all()
172
+
173
+ root = (2 * c) / (-b - torch.sqrt(discriminant))
174
+ outputs = root * input_bin_widths + input_cumwidths
175
+
176
+ theta_one_minus_theta = root * (1 - root)
177
+ denominator = input_delta + (
178
+ (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
179
+ * theta_one_minus_theta
180
+ )
181
+ derivative_numerator = input_delta.pow(2) * (
182
+ input_derivatives_plus_one * root.pow(2)
183
+ + 2 * input_delta * theta_one_minus_theta
184
+ + input_derivatives * (1 - root).pow(2)
185
+ )
186
+ logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
187
+
188
+ return outputs, -logabsdet
189
+ else:
190
+ theta = (inputs - input_cumwidths) / input_bin_widths
191
+ theta_one_minus_theta = theta * (1 - theta)
192
+
193
+ numerator = input_heights * (
194
+ input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
195
+ )
196
+ denominator = input_delta + (
197
+ (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
198
+ * theta_one_minus_theta
199
+ )
200
+ outputs = input_cumheights + numerator / denominator
201
+
202
+ derivative_numerator = input_delta.pow(2) * (
203
+ input_derivatives_plus_one * theta.pow(2)
204
+ + 2 * input_delta * theta_one_minus_theta
205
+ + input_derivatives * (1 - theta).pow(2)
206
+ )
207
+ logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
208
+
209
+ return outputs, logabsdet
lipsync-models/GPT_SoVITS/prepare_datasets/1-get-text.py ADDED
@@ -0,0 +1,131 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import os
4
+
5
+ inp_text = os.environ.get("inp_text")
6
+ inp_wav_dir = os.environ.get("inp_wav_dir")
7
+ exp_name = os.environ.get("exp_name")
8
+ i_part = os.environ.get("i_part")
9
+ all_parts = os.environ.get("all_parts")
10
+ os.environ["CUDA_VISIBLE_DEVICES"] = os.environ.get("_CUDA_VISIBLE_DEVICES")
11
+ opt_dir = os.environ.get("opt_dir")
12
+ bert_pretrained_dir = os.environ.get("bert_pretrained_dir")
13
+ is_half = eval(os.environ.get("is_half", "True"))
14
+ import sys, numpy as np, traceback, pdb
15
+ import os.path
16
+ from glob import glob
17
+ from tqdm import tqdm
18
+ from text.cleaner import clean_text
19
+ import torch
20
+ from transformers import AutoModelForMaskedLM, AutoTokenizer
21
+ import numpy as np
22
+
23
+ # inp_text=sys.argv[1]
24
+ # inp_wav_dir=sys.argv[2]
25
+ # exp_name=sys.argv[3]
26
+ # i_part=sys.argv[4]
27
+ # all_parts=sys.argv[5]
28
+ # os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[6]#i_gpu
29
+ # opt_dir="/data/docker/liujing04/gpt-vits/fine_tune_dataset/%s"%exp_name
30
+ # bert_pretrained_dir="/data/docker/liujing04/bert-vits2/Bert-VITS2-master20231106/bert/chinese-roberta-wwm-ext-large"
31
+
32
+ from time import time as ttime
33
+ import shutil
34
+
35
+
36
+ def my_save(fea,path):#####fix issue: torch.save doesn't support chinese path
37
+ dir=os.path.dirname(path)
38
+ name=os.path.basename(path)
39
+ # tmp_path="%s/%s%s.pth"%(dir,ttime(),i_part)
40
+ tmp_path="%s%s.pth"%(ttime(),i_part)
41
+ torch.save(fea,tmp_path)
42
+ shutil.move(tmp_path,"%s/%s"%(dir,name))
43
+
44
+
45
+ txt_path = "%s/2-name2text-%s.txt" % (opt_dir, i_part)
46
+ if os.path.exists(txt_path) == False:
47
+ bert_dir = "%s/3-bert" % (opt_dir)
48
+ os.makedirs(opt_dir, exist_ok=True)
49
+ os.makedirs(bert_dir, exist_ok=True)
50
+ if torch.cuda.is_available():
51
+ device = "cuda:0"
52
+ elif torch.backends.mps.is_available():
53
+ device = "mps"
54
+ else:
55
+ device = "cpu"
56
+ tokenizer = AutoTokenizer.from_pretrained(bert_pretrained_dir)
57
+ bert_model = AutoModelForMaskedLM.from_pretrained(bert_pretrained_dir)
58
+ if is_half == True:
59
+ bert_model = bert_model.half().to(device)
60
+ else:
61
+ bert_model = bert_model.to(device)
62
+
63
+ def get_bert_feature(text, word2ph):
64
+ with torch.no_grad():
65
+ inputs = tokenizer(text, return_tensors="pt")
66
+ for i in inputs:
67
+ inputs[i] = inputs[i].to(device)
68
+ res = bert_model(**inputs, output_hidden_states=True)
69
+ res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1]
70
+
71
+ assert len(word2ph) == len(text)
72
+ phone_level_feature = []
73
+ for i in range(len(word2ph)):
74
+ repeat_feature = res[i].repeat(word2ph[i], 1)
75
+ phone_level_feature.append(repeat_feature)
76
+
77
+ phone_level_feature = torch.cat(phone_level_feature, dim=0)
78
+
79
+ return phone_level_feature.T
80
+
81
+ def process(data, res):
82
+ for name, text, lan in data:
83
+ try:
84
+ name = os.path.basename(name)
85
+ phones, word2ph, norm_text = clean_text(
86
+ text.replace("%", "-").replace("¥", ","), lan
87
+ )
88
+ path_bert = "%s/%s.pt" % (bert_dir, name)
89
+ if os.path.exists(path_bert) == False and lan == "zh":
90
+ bert_feature = get_bert_feature(norm_text, word2ph)
91
+ assert bert_feature.shape[-1] == len(phones)
92
+ # torch.save(bert_feature, path_bert)
93
+ my_save(bert_feature, path_bert)
94
+ phones = " ".join(phones)
95
+ # res.append([name,phones])
96
+ res.append([name, phones, word2ph, norm_text])
97
+ except:
98
+ print(name, text, traceback.format_exc())
99
+
100
+ todo = []
101
+ res = []
102
+ with open(inp_text, "r", encoding="utf8") as f:
103
+ lines = f.read().strip("\n").split("\n")
104
+
105
+ language_v1_to_language_v2 = {
106
+ "ZH": "zh",
107
+ "zh": "zh",
108
+ "JP": "ja",
109
+ "jp": "ja",
110
+ "JA": "ja",
111
+ "ja": "ja",
112
+ "EN": "en",
113
+ "en": "en",
114
+ "En": "en",
115
+ }
116
+ for line in lines[int(i_part) :: int(all_parts)]:
117
+ try:
118
+ wav_name, spk_name, language, text = line.split("|")
119
+ # todo.append([name,text,"zh"])
120
+ todo.append(
121
+ [wav_name, text, language_v1_to_language_v2.get(language, language)]
122
+ )
123
+ except:
124
+ print(line, traceback.format_exc())
125
+
126
+ process(todo, res)
127
+ opt = []
128
+ for name, phones, word2ph, norm_text in res:
129
+ opt.append("%s\t%s\t%s\t%s" % (name, phones, word2ph, norm_text))
130
+ with open(txt_path, "w", encoding="utf8") as f:
131
+ f.write("\n".join(opt) + "\n")