55	
56	    def check_values(self):
57	        super().check_values()
58	        c = asdict(self)
59	        assert (
60	            c["model_params"]["input_dim"] == self.audio.num_mels
61	        ), " [!] model input dimendion must be equal to melspectrogram dimension."

125	            # get classid
126	            class_id = self.classname_to_classid[class_name]
127	            # load wav file
128	            wav = self.load_wav(utter_path)
129	            offset = random.randint(0, wav.shape[0] - self.seq_len)
130	            wav = wav[offset : offset + self.seq_len]
131	
132	            if self.augmentator is not None and self.data_augmentation_p:

129	            offset = random.randint(0, wav.shape[0] - self.seq_len)
130	            wav = wav[offset : offset + self.seq_len]
131	
132	            if self.augmentator is not None and self.data_augmentation_p:
133	                if random.random() < self.data_augmentation_p:
134	                    wav = self.augmentator.apply_one(wav)
135	
136	            if not self.use_torch_spec:

24	        self.loss_method = loss_method
25	
26	        print(" > Initialized Generalized End-to-End loss")
27	
28	        assert self.loss_method in ["softmax", "contrast"]
29	
30	        if self.loss_method == "softmax":
31	            self.embed_loss = self.embed_loss_softmax

107	        """
108	        Calculates the GE2E loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats)
109	        """
110	
111	        assert x.size()[1] >= 2
112	
113	        centroids = torch.mean(x, 1)
114	        cos_sim_matrix = self.calc_cosine_sim(x, centroids)

145	        """
146	        Calculates the AngleProto loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats)
147	        """
148	
149	        assert x.size()[1] >= 2
150	
151	        out_anchor = torch.mean(x[:, 1:, :], 1)
152	        out_positive = x[:, 0, :]

15	        self.coefficient = coefficient
16	        self.register_buffer("filter", torch.FloatTensor([-self.coefficient, 1.0]).unsqueeze(0).unsqueeze(0))
17	
18	    def forward(self, x):
19	        assert len(x.size()) == 2
20	
21	        x = torch.nn.functional.pad(x.unsqueeze(1), (1, 0), "reflect")
22	        return torch.nn.functional.conv1d(x, self.filter).squeeze(1)

153	                criterion = criterion.cuda()
154	
155	        if eval:
156	            self.eval()
157	            assert not self.training
158	
159	        if not eval:
160	            return criterion, state["step"]

66	        clean_db = 10 * np.log10(np.mean(audio**2) + 1e-4)
67	
68	        noise_list = random.sample(
69	            self.noise_list[noise_type],
70	            random.randint(
71	                self.additive_noise_config[noise_type]["min_num_noises"],
72	                self.additive_noise_config[noise_type]["max_num_noises"],
73	            ),
74	        )
75	
76	        audio_len = audio.shape[0]

80	
81	            if noiseaudio.shape[0] < audio_len:
82	                continue
83	
84	            noise_snr = random.uniform(
85	                self.additive_noise_config[noise_type]["min_snr_in_db"],
86	                self.additive_noise_config[noise_type]["max_num_noises"],
87	            )
88	            noise_db = 10 * np.log10(np.mean(noiseaudio**2) + 1e-4)
89	            noise_wav = np.sqrt(10 ** ((clean_db - noise_db - noise_snr) / 10)) * noiseaudio
90	

101	
102	    def reverberate(self, audio):
103	        audio_len = audio.shape[0]
104	
105	        rir_file = random.choice(self.rir_files)
106	        rir = self.ap.load_wav(rir_file, sr=self.ap.sample_rate)
107	        rir = rir / np.sqrt(np.sum(rir**2))
108	        return signal.convolve(audio, rir, mode=self.rir_config["conv_mode"])[:audio_len]

107	        rir = rir / np.sqrt(np.sum(rir**2))
108	        return signal.convolve(audio, rir, mode=self.rir_config["conv_mode"])[:audio_len]
109	
110	    def apply_one(self, audio):
111	        noise_type = random.choice(self.global_noise_list)
112	        if noise_type == "RIR_AUG":
113	            return self.reverberate(audio)
114	

20	""" voxceleb 1 & 2 """
21	
22	import hashlib
23	import os
24	import subprocess
25	import sys
26	import zipfile
27	

160	        Raises:
161	            AssertionError: when the parameters network is not defined
162	            AssertionError: transition probability is not between 0 and 1
163	        """
164	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
165	        assert (
166	            len(self.outputnet_size) >= 1
167	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"

161	            AssertionError: when the parameters network is not defined
162	            AssertionError: transition probability is not between 0 and 1
163	        """
164	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
165	        assert (
166	            len(self.outputnet_size) >= 1
167	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
168	        assert (
169	            0 < self.flat_start_params["transition_p"] < 1
170	        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"

164	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
165	        assert (
166	            len(self.outputnet_size) >= 1
167	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
168	        assert (
169	            0 < self.flat_start_params["transition_p"] < 1
170	        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"

191	        Raises:
192	            AssertionError: when the parameters network is not defined
193	            AssertionError: transition probability is not between 0 and 1
194	        """
195	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
196	        assert (
197	            len(self.outputnet_size) >= 1
198	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"

192	            AssertionError: when the parameters network is not defined
193	            AssertionError: transition probability is not between 0 and 1
194	        """
195	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
196	        assert (
197	            len(self.outputnet_size) >= 1
198	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
199	        assert (
200	            0 < self.flat_start_params["transition_p"] < 1
201	        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"

195	        assert self.ar_order > 0, "AR order must be greater than 0 it is an autoregressive model."
196	        assert (
197	            len(self.outputnet_size) >= 1
198	        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
199	        assert (
200	            0 < self.flat_start_params["transition_p"] < 1
201	        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"

223	    )
224	
225	    def check_values(self):
226	        if self.gradual_training:
227	            assert (
228	                self.gradual_training[0][1] == self.r
229	            ), f"[!] the first scheduled gradual training `r` must be equal to the model's `r` value. {self.gradual_training[0][1]} vs {self.r}"
230	        if self.model == "tacotron" and self.audio is not None:
231	            assert self.out_channels == (
232	                self.audio.fft_size // 2 + 1

227	            assert (
228	                self.gradual_training[0][1] == self.r
229	            ), f"[!] the first scheduled gradual training `r` must be equal to the model's `r` value. {self.gradual_training[0][1]} vs {self.r}"
230	        if self.model == "tacotron" and self.audio is not None:
231	            assert self.out_channels == (
232	                self.audio.fft_size // 2 + 1
233	            ), f"{self.out_channels} vs {self.audio.fft_size // 2 + 1}"
234	        if self.model == "tacotron2" and self.audio is not None:
235	            assert self.out_channels == self.audio.num_mels

231	            assert self.out_channels == (
232	                self.audio.fft_size // 2 + 1
233	            ), f"{self.out_channels} vs {self.audio.fft_size // 2 + 1}"
234	        if self.model == "tacotron2" and self.audio is not None:
235	            assert self.out_channels == self.audio.num_mels

33	            eval_split_size = min(eval_split_max_size, int(len(items) * eval_split_size))
34	        else:
35	            eval_split_size = int(len(items) * eval_split_size)
36	
37	    assert (
38	        eval_split_size > 0
39	    ), " [!] You do not have enough samples for the evaluation set. You can work around this setting the 'eval_split_size' parameter to a minimum of {}".format(
40	        1 / len(items)
41	    )
42	    np.random.seed(0)
43	    np.random.shuffle(items)
44	    if is_multi_speaker:

117	        if formatter is None:
118	            formatter = _get_formatter_by_name(formatter_name)
119	        # load train set
120	        meta_data_train = formatter(root_path, meta_file_train, ignored_speakers=ignored_speakers)
121	        assert len(meta_data_train) > 0, f" [!] No training samples found in {root_path}/{meta_file_train}"
122	
123	        meta_data_train = add_extra_keys(meta_data_train, language, dataset_name)
124	

220	        print(f"{indent}| > Number of instances : {len(self.samples)}")
221	
222	    def load_wav(self, filename):
223	        waveform = self.ap.load_wav(filename)
224	        assert waveform.size > 0
225	        return waveform
226	
227	    def get_phonemes(self, idx, text):

225	        return waveform
226	
227	    def get_phonemes(self, idx, text):
228	        out_dict = self.phoneme_dataset[idx]
229	        assert text == out_dict["text"], f"{text} != {out_dict['text']}"
230	        assert len(out_dict["token_ids"]) > 0
231	        return out_dict
232	

226	
227	    def get_phonemes(self, idx, text):
228	        out_dict = self.phoneme_dataset[idx]
229	        assert text == out_dict["text"], f"{text} != {out_dict['text']}"
230	        assert len(out_dict["token_ids"]) > 0
231	        return out_dict
232	
233	    def get_f0(self, idx):

232	
233	    def get_f0(self, idx):
234	        out_dict = self.f0_dataset[idx]
235	        item = self.samples[idx]
236	        assert item["audio_unique_name"] == out_dict["audio_unique_name"]
237	        return out_dict
238	
239	    def get_energy(self, idx):

238	
239	    def get_energy(self, idx):
240	        out_dict = self.energy_dataset[idx]
241	        item = self.samples[idx]
242	        assert item["audio_unique_name"] == out_dict["audio_unique_name"]
243	        return out_dict
244	
245	    @staticmethod

333	        return idxs
334	
335	    @staticmethod
336	    def create_buckets(samples, batch_group_size: int):
337	        assert batch_group_size > 0
338	        for i in range(len(samples) // batch_group_size):
339	            offset = i * batch_group_size
340	            end_offset = offset + batch_group_size

499	            if self.compute_linear_spec:
500	                linear = [self.ap.spectrogram(w).astype("float32") for w in batch["wav"]]
501	                linear = prepare_tensor(linear, self.outputs_per_step)
502	                linear = linear.transpose(0, 2, 1)
503	                assert mel.shape[1] == linear.shape[1]
504	                linear = torch.FloatTensor(linear).contiguous()
505	
506	            # format waveforms

519	
520	            # format F0
521	            if self.compute_f0:
522	                pitch = prepare_data(batch["pitch"])
523	                assert mel.shape[1] == pitch.shape[1], f"[!] {mel.shape} vs {pitch.shape}"
524	                pitch = torch.FloatTensor(pitch)[:, None, :].contiguous()  # B x 1 xT
525	            else:
526	                pitch = None

526	                pitch = None
527	            # format energy
528	            if self.compute_energy:
529	                energy = prepare_data(batch["energy"])
530	                assert mel.shape[1] == energy.shape[1], f"[!] {mel.shape} vs {energy.shape}"
531	                energy = torch.FloatTensor(energy)[:, None, :].contiguous()  # B x 1 xT
532	            else:
533	                energy = None

537	                attns = [batch["attn"][idx].T for idx in ids_sorted_decreasing]
538	                for idx, attn in enumerate(attns):
539	                    pad2 = mel.shape[1] - attn.shape[1]
540	                    pad1 = token_ids.shape[1] - attn.shape[0]
541	                    assert pad1 >= 0 and pad2 >= 0, f"[!] Negative padding - {pad1} and {pad2}"
542	                    attn = np.pad(attn, [[0, pad1], [0, pad2]])
543	                    attns[idx] = attn
544	                attns = prepare_tensor(attns, self.outputs_per_step)

720	    def __getitem__(self, idx):
721	        item = self.samples[idx]
722	        f0 = self.compute_or_load(item["audio_file"], string2filename(item["audio_unique_name"]))
723	        if self.normalize_f0:
724	            assert self.mean is not None and self.std is not None, " [!] Mean and STD is not available"
725	            f0 = self.normalize(f0)
726	        return {"audio_unique_name": item["audio_unique_name"], "f0": f0}
727	

871	    def __getitem__(self, idx):
872	        item = self.samples[idx]
873	        energy = self.compute_or_load(item["audio_file"], string2filename(item["audio_unique_name"]))
874	        if self.normalize_energy:
875	            assert self.mean is not None and self.std is not None, " [!] Mean and STD is not available"
876	            energy = self.normalize(energy)
877	        return {"audio_unique_name": item["audio_unique_name"], "energy": energy}
878	

1	import os
2	import re
3	import xml.etree.ElementTree as ET
4	from glob import glob
5	from pathlib import Path
6	from typing import List
7	
8	import pandas as pd

25	        if len(line.split("|")) != num_cols:
26	            print(f" > Missing column in line {idx + 1} -> {line.strip()}")
27	    # load metadata
28	    metadata = pd.read_csv(os.path.join(root_path, meta_file), sep="|")
29	    assert all(x in metadata.columns for x in ["wav_filename", "transcript"])
30	    client_id = None if "client_id" in metadata.columns else "default"
31	    emotion_name = None if "emotion_name" in metadata.columns else "neutral"
32	    items = []

63	        if len(line.split("|")) != num_cols:
64	            print(f" > Missing column in line {idx + 1} -> {line.strip()}")
65	    # load metadata
66	    metadata = pd.read_csv(os.path.join(root_path, meta_file), sep="|")
67	    assert all(x in metadata.columns for x in ["audio_file", "text"])
68	    speaker_name = None if "speaker_name" in metadata.columns else "coqui"
69	    emotion_name = None if "emotion_name" in metadata.columns else "neutral"
70	    items = []

351	                        "root_path": root_path,
352	                    }
353	                )
354	    for item in items:
355	        assert os.path.exists(item["audio_file"]), f" [!] wav files don't exist - {item['audio_file']}"
356	    return items
357	
358	

537	    return _voxcel_x(root_path, meta_file, voxcel_idx="1")
538	
539	
540	def _voxcel_x(root_path, meta_file, voxcel_idx):
541	    assert voxcel_idx in ["1", "2"]
542	    expected_count = 148_000 if voxcel_idx == "1" else 1_000_000
543	    voxceleb_path = Path(root_path)
544	    cache_to = voxceleb_path / f"metafile_voxceleb{voxcel_idx}.csv"

558	            desc=f"Building VoxCeleb {voxcel_idx} Meta file ... this needs to be done only once.",
559	            total=expected_count,
560	        ):
561	            speaker_id = str(Path(path).parent.parent.stem)
562	            assert speaker_id.startswith("id")
563	            text = None  # VoxCel does not provide transciptions, and they are not needed for training the SE
564	            meta_data.append(f"{text}|{path}|voxcel{voxcel_idx}_{speaker_id}\n")
565	            cnt += 1

183	
184	    Returns:
185	        np.ndarray: The generated semantic tokens.
186	    """
187	    assert isinstance(text, str)
188	    text = _normalize_whitespace(text)
189	    assert len(text.strip()) > 0
190	    if all(v is not None for v in history_prompt) or base is not None:

185	        np.ndarray: The generated semantic tokens.
186	    """
187	    assert isinstance(text, str)
188	    text = _normalize_whitespace(text)
189	    assert len(text.strip()) > 0
190	    if all(v is not None for v in history_prompt) or base is not None:
191	        if history_prompt is not None:
192	            semantic_history = history_prompt[0]

191	        if history_prompt is not None:
192	            semantic_history = history_prompt[0]
193	        if base is not None:
194	            semantic_history = base[0]
195	        assert (
196	            isinstance(semantic_history, np.ndarray)
197	            and len(semantic_history.shape) == 1
198	            and len(semantic_history) > 0
199	            and semantic_history.min() >= 0
200	            and semantic_history.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
201	        )
202	    else:
203	        semantic_history = None
204	    encoded_text = np.array(_tokenize(model.tokenizer, text)) + model.config.TEXT_ENCODING_OFFSET

226	        semantic_history = np.array([model.config.SEMANTIC_PAD_TOKEN] * 256)
227	    x = torch.from_numpy(
228	        np.hstack([encoded_text, semantic_history, np.array([model.config.SEMANTIC_INFER_TOKEN])]).astype(np.int64)
229	    )[None]
230	    assert x.shape[1] == 256 + 256 + 1
231	    with inference_mode():
232	        x = x.to(model.device)
233	        n_tot_steps = 768

288	                pbar.update(req_pbar_state - pbar_state)
289	            pbar_state = req_pbar_state
290	        pbar.close()
291	        out = x.detach().cpu().numpy().squeeze()[256 + 256 + 1 :]
292	    assert all(out >= 0) and all(out < model.config.SEMANTIC_VOCAB_SIZE)
293	    clear_cuda_cache()
294	    return out
295	

294	    return out
295	
296	
297	def _flatten_codebooks(arr, offset_size):
298	    assert len(arr.shape) == 2
299	    arr = arr.copy()
300	    if offset_size is not None:
301	        for n in range(1, arr.shape[0]):

334	
335	    Returns:
336	        np.ndarray: The generated coarse audio codes.
337	    """
338	    assert (
339	        isinstance(x_semantic, np.ndarray)
340	        and len(x_semantic.shape) == 1
341	        and len(x_semantic) > 0
342	        and x_semantic.min() >= 0
343	        and x_semantic.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
344	    )
345	    assert 60 <= max_coarse_history <= 630
346	    assert max_coarse_history + sliding_window_len <= 1024 - 256
347	    semantic_to_coarse_ratio = (

341	        and len(x_semantic) > 0
342	        and x_semantic.min() >= 0
343	        and x_semantic.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
344	    )
345	    assert 60 <= max_coarse_history <= 630
346	    assert max_coarse_history + sliding_window_len <= 1024 - 256
347	    semantic_to_coarse_ratio = (
348	        model.config.COARSE_RATE_HZ / model.config.SEMANTIC_RATE_HZ * model.config.N_COARSE_CODEBOOKS

342	        and x_semantic.min() >= 0
343	        and x_semantic.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
344	    )
345	    assert 60 <= max_coarse_history <= 630
346	    assert max_coarse_history + sliding_window_len <= 1024 - 256
347	    semantic_to_coarse_ratio = (
348	        model.config.COARSE_RATE_HZ / model.config.SEMANTIC_RATE_HZ * model.config.N_COARSE_CODEBOOKS
349	    )

355	            x_coarse_history = x_history[1]
356	        if base is not None:
357	            x_semantic_history = base[0]
358	            x_coarse_history = base[1]
359	        assert (
360	            isinstance(x_semantic_history, np.ndarray)
361	            and len(x_semantic_history.shape) == 1
362	            and len(x_semantic_history) > 0
363	            and x_semantic_history.min() >= 0
364	            and x_semantic_history.max() <= model.config.SEMANTIC_VOCAB_SIZE - 1
365	            and isinstance(x_coarse_history, np.ndarray)
366	            and len(x_coarse_history.shape) == 2
367	            and x_coarse_history.shape[0] == model.config.N_COARSE_CODEBOOKS
368	            and x_coarse_history.shape[-1] >= 0
369	            and x_coarse_history.min() >= 0
370	            and x_coarse_history.max() <= model.config.CODEBOOK_SIZE - 1
371	            and (
372	                round(x_coarse_history.shape[-1] / len(x_semantic_history), 1)
373	                == round(semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS, 1)
374	            )
375	        )
376	        x_coarse_history = (
377	            _flatten_codebooks(x_coarse_history, model.config.CODEBOOK_SIZE) + model.config.SEMANTIC_VOCAB_SIZE
378	        )

398	            np.floor(len(x_semantic) * semantic_to_coarse_ratio / model.config.N_COARSE_CODEBOOKS)
399	            * model.config.N_COARSE_CODEBOOKS
400	        )
401	    )
402	    assert n_steps > 0 and n_steps % model.config.N_COARSE_CODEBOOKS == 0
403	    x_semantic = np.hstack([x_semantic_history, x_semantic]).astype(np.int32)
404	    x_coarse = x_coarse_history.astype(np.int32)
405	    base_semantic_idx = len(x_semantic_history)

470	            del x_in
471	        del x_semantic_in
472	    gen_coarse_arr = x_coarse_in.detach().cpu().numpy().squeeze()[len(x_coarse_history) :]
473	    del x_coarse_in
474	    assert len(gen_coarse_arr) == n_steps
475	    gen_coarse_audio_arr = (
476	        gen_coarse_arr.reshape(-1, model.config.N_COARSE_CODEBOOKS).T - model.config.SEMANTIC_VOCAB_SIZE
477	    )

501	
502	    Returns:
503	        np.ndarray: The generated full audio codes.
504	    """
505	    assert (
506	        isinstance(x_coarse_gen, np.ndarray)
507	        and len(x_coarse_gen.shape) == 2
508	        and 1 <= x_coarse_gen.shape[0] <= model.config.N_FINE_CODEBOOKS - 1
509	        and x_coarse_gen.shape[1] > 0
510	        and x_coarse_gen.min() >= 0
511	        and x_coarse_gen.max() <= model.config.CODEBOOK_SIZE - 1
512	    )
513	    if all(v is not None for v in history_prompt) or base is not None:
514	        if history_prompt is not None:
515	            x_fine_history = history_prompt[2]

514	        if history_prompt is not None:
515	            x_fine_history = history_prompt[2]
516	        if base is not None:
517	            x_fine_history = base[2]
518	        assert (
519	            isinstance(x_fine_history, np.ndarray)
520	            and len(x_fine_history.shape) == 2
521	            and x_fine_history.shape[0] == model.config.N_FINE_CODEBOOKS
522	            and x_fine_history.shape[1] >= 0
523	            and x_fine_history.min() >= 0
524	            and x_fine_history.max() <= model.config.CODEBOOK_SIZE - 1
525	        )
526	    else:
527	        x_fine_history = None
528	    n_coarse = x_coarse_gen.shape[0]

589	        del in_arr
590	    gen_fine_arr = gen_fine_arr[:, n_history:]
591	    if n_remove_from_end > 0:
592	        gen_fine_arr = gen_fine_arr[:, :-n_remove_from_end]
593	    assert gen_fine_arr.shape[-1] == x_coarse_gen.shape[-1]
594	    clear_cuda_cache()
595	    return gen_fine_arr
596	

25	
26	class CausalSelfAttention(nn.Module):
27	    def __init__(self, config):
28	        super().__init__()
29	        assert config.n_embd % config.n_head == 0
30	        # key, query, value projections for all heads, but in a batch
31	        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
32	        # output projection

145	
146	class GPT(nn.Module):
147	    def __init__(self, config):
148	        super().__init__()
149	        assert config.input_vocab_size is not None
150	        assert config.output_vocab_size is not None
151	        assert config.block_size is not None
152	        self.config = config

146	class GPT(nn.Module):
147	    def __init__(self, config):
148	        super().__init__()
149	        assert config.input_vocab_size is not None
150	        assert config.output_vocab_size is not None
151	        assert config.block_size is not None
152	        self.config = config
153	

147	    def __init__(self, config):
148	        super().__init__()
149	        assert config.input_vocab_size is not None
150	        assert config.output_vocab_size is not None
151	        assert config.block_size is not None
152	        self.config = config
153	
154	        self.transformer = nn.ModuleDict(

178	    def forward(self, idx, merge_context=False, past_kv=None, position_ids=None, use_cache=False):
179	        device = idx.device
180	        _, t = idx.size()
181	        if past_kv is not None:
182	            assert t == 1
183	            tok_emb = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
184	        else:
185	            if merge_context:

182	            assert t == 1
183	            tok_emb = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
184	        else:
185	            if merge_context:
186	                assert idx.shape[1] >= 256 + 256 + 1
187	                t = idx.shape[1] - 256
188	            else:
189	                assert (

185	            if merge_context:
186	                assert idx.shape[1] >= 256 + 256 + 1
187	                t = idx.shape[1] - 256
188	            else:
189	                assert (
190	                    t <= self.config.block_size
191	                ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
192	
193	            # forward the GPT model itself
194	            if merge_context:

210	
211	        if position_ids is None:
212	            position_ids = torch.arange(past_length, t + past_length, dtype=torch.long, device=device)
213	            position_ids = position_ids.unsqueeze(0)  # shape (1, t)
214	            assert position_ids.shape == (1, t)
215	
216	        pos_emb = self.transformer.wpe(position_ids)  # position embeddings of shape (1, t, n_embd)
217	

14	
15	class NonCausalSelfAttention(nn.Module):
16	    def __init__(self, config):
17	        super().__init__()
18	        assert config.n_embd % config.n_head == 0
19	        # key, query, value projections for all heads, but in a batch
20	        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
21	        # output projection

99	
100	    def forward(self, pred_idx, idx):
101	        device = idx.device
102	        b, t, codes = idx.size()
103	        assert (
104	            t <= self.config.block_size
105	        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
106	        assert pred_idx > 0, "cannot predict 0th codebook"
107	        assert codes == self.n_codes_total, (b, t, codes)
108	        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)

102	        b, t, codes = idx.size()
103	        assert (
104	            t <= self.config.block_size
105	        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
106	        assert pred_idx > 0, "cannot predict 0th codebook"
107	        assert codes == self.n_codes_total, (b, t, codes)
108	        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
109	

103	        assert (
104	            t <= self.config.block_size
105	        ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
106	        assert pred_idx > 0, "cannot predict 0th codebook"
107	        assert codes == self.n_codes_total, (b, t, codes)
108	        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # shape (1, t)
109	
110	        # forward the GPT model itself

353	        d_model: int = 512,
354	        num_heads: int = 16,
355	    ):
356	        super().__init__()
357	        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
358	        self.d_model = d_model
359	        self.d_head = int(d_model / num_heads)
360	        self.num_heads = num_heads

56	        use_weight_norm=False,
57	    ):
58	        super(ConvNorm, self).__init__()  # pylint: disable=super-with-arguments
59	        if padding is None:
60	            assert kernel_size % 2 == 1
61	            padding = int(dilation * (kernel_size - 1) / 2)
62	        self.kernel_size = kernel_size
63	        self.dilation = dilation

640	            (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length).
641	        """
642	        batch, _, in_length = x.shape
643	        batch, _, out_channels, kernel_size, kernel_length = kernel.shape
644	        assert in_length == (kernel_length * hop_size), "length of (x, kernel) is not matched"
645	
646	        padding = dilation * int((kernel_size - 1) / 2)
647	        x = F.pad(x, (padding, padding), "constant", 0)  # (batch, in_channels, in_length + 2*padding)

9	from TTS.tts.layers.delightful_tts.conv_layers import ConvNorm
10	
11	
12	def initialize_embeddings(shape: Tuple[int]) -> torch.Tensor:
13	    assert len(shape) == 2, "Can only initialize 2-D embedding matrices ..."
14	    # Kaiming initialization
15	    return torch.randn(shape) * np.sqrt(2 / shape[1])
16	

142	            )
143	        elif encoder_type.lower() == "residual_conv_bn":
144	            self.encoder = ResidualConv1dBNEncoder(in_hidden_channels, out_channels, in_hidden_channels, encoder_params)
145	        elif encoder_type.lower() == "fftransformer":
146	            assert (
147	                in_hidden_channels == out_channels
148	            ), "[!] must be `in_channels` == `out_channels` when encoder type is 'fftransformer'"
149	            # pylint: disable=unexpected-keyword-arg
150	            self.encoder = FFTransformerBlock(in_hidden_channels, **encoder_params)
151	        else:

100	    def __init__(
101	        self, in_channels, out_channels, hidden_channels, kernel_size, dilations, num_res_blocks=13, num_conv_blocks=2
102	    ):
103	        super().__init__()
104	        assert len(dilations) == num_res_blocks
105	        self.res_blocks = nn.ModuleList()
106	        for idx, dilation in enumerate(dilations):
107	            block = Conv1dBNBlock(

59	
60	class TimeDepthSeparableConvBlock(nn.Module):
61	    def __init__(self, in_channels, hid_channels, out_channels, num_layers, kernel_size, bias=True):
62	        super().__init__()
63	        assert (kernel_size - 1) % 2 == 0
64	        assert num_layers > 1
65	
66	        self.layers = nn.ModuleList()

60	class TimeDepthSeparableConvBlock(nn.Module):
61	    def __init__(self, in_channels, hid_channels, out_channels, num_layers, kernel_size, bias=True):
62	        super().__init__()
63	        assert (kernel_size - 1) % 2 == 0
64	        assert num_layers > 1
65	
66	        self.layers = nn.ModuleList()
67	        layer = TimeDepthSeparableConv(

45	        dropout_p=0,
46	        weight_norm=True,
47	    ):
48	        super().__init__()
49	        assert kernel_size % 2 == 1
50	        assert hidden_channels % 2 == 0
51	        self.in_channels = in_channels
52	        self.hidden_channels = hidden_channels

46	        weight_norm=True,
47	    ):
48	        super().__init__()
49	        assert kernel_size % 2 == 1
50	        assert hidden_channels % 2 == 0
51	        self.in_channels = in_channels
52	        self.hidden_channels = hidden_channels
53	        self.kernel_size = kernel_size

33	        self.out_channels = out_channels
34	        self.kernel_size = kernel_size
35	        self.num_layers = num_layers
36	        self.dropout_p = dropout_p
37	        assert num_layers > 1, " [!] number of layers should be > 0."
38	        assert kernel_size % 2 == 1, " [!] kernel size should be odd number."
39	
40	        self.conv_layers = nn.ModuleList()

34	        self.kernel_size = kernel_size
35	        self.num_layers = num_layers
36	        self.dropout_p = dropout_p
37	        assert num_layers > 1, " [!] number of layers should be > 0."
38	        assert kernel_size % 2 == 1, " [!] kernel size should be odd number."
39	
40	        self.conv_layers = nn.ModuleList()
41	        self.norm_layers = nn.ModuleList()

83	    """
84	
85	    def __init__(self, channels, num_splits=4, no_jacobian=False, **kwargs):  # pylint: disable=unused-argument
86	        super().__init__()
87	        assert num_splits % 2 == 0
88	        self.channels = channels
89	        self.num_splits = num_splits
90	        self.no_jacobian = no_jacobian

105	            - x: :math:`[B, C, T]`
106	            - x_mask: :math:`[B, 1, T]`
107	        """
108	        b, c, t = x.size()
109	        assert c % self.num_splits == 0
110	        if x_mask is None:
111	            x_mask = 1
112	            x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t

64	        proximal_bias=False,
65	        proximal_init=False,
66	    ):
67	        super().__init__()
68	        assert channels % num_heads == 0, " [!] channels should be divisible by num_heads."
69	        # class attributes
70	        self.channels = channels
71	        self.out_channels = out_channels

129	        # compute raw attention scores
130	        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels)
131	        # relative positional encoding for scores
132	        if self.rel_attn_window_size is not None:
133	            assert t_s == t_t, "Relative attention is only available for self-attention."
134	            # get relative key embeddings
135	            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
136	            rel_logits = self._matmul_with_relative_keys(query, key_relative_embeddings)

138	            scores_local = rel_logits / math.sqrt(self.k_channels)
139	            scores = scores + scores_local
140	        # proximan bias
141	        if self.proximal_bias:
142	            assert t_s == t_t, "Proximal bias is only available for self-attention."
143	            scores = scores + self._attn_proximity_bias(t_s).to(device=scores.device, dtype=scores.dtype)
144	        # attention score masking
145	        if mask is not None:

73	
74	        self.transition_model = TransitionModel()
75	        self.emission_model = EmissionModel()
76	
77	        assert ar_order > 0, f"AR order must be greater than 0 provided {ar_order}"
78	
79	        self.ar_order = ar_order
80	        self.prenet = Prenet(

168	        for conv1d in self.conv1d_banks:
169	            out = conv1d(x)
170	            outs.append(out)
171	        x = torch.cat(outs, dim=1)
172	        assert x.size(1) == self.conv_bank_features * len(self.conv1d_banks)
173	        for conv1d in self.conv1d_projections:
174	            x = conv1d(x)
175	        x += inputs

23	    """
24	
25	    def __init__(self, in_channels, out_channels, kernel_size, activation=None):
26	        super().__init__()
27	        assert (kernel_size - 1) % 2 == 0
28	        padding = (kernel_size - 1) // 2
29	        self.convolution1d = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding)
30	        self.batch_normalization = nn.BatchNorm1d(out_channels, momentum=0.1, eps=1e-5)

39	    elif channels <= 64:
40	        groups = 16
41	    while channels % groups != 0:
42	        groups = int(groups / 2)
43	    assert groups > 2
44	    return GroupNorm32(groups, channels)
45	
46	

60	        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
61	        :return: an [N x (H * C) x T] tensor after attention.
62	        """
63	        bs, width, length = qkv.shape
64	        assert width % (3 * self.n_heads) == 0
65	        ch = width // (3 * self.n_heads)
66	        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
67	        scale = 1 / math.sqrt(math.sqrt(ch))

101	        self.do_checkpoint = do_checkpoint
102	        if num_head_channels == -1:
103	            self.num_heads = num_heads
104	        else:
105	            assert (
106	                channels % num_head_channels == 0
107	            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
108	            self.num_heads = channels // num_head_channels
109	        self.norm = normalization(channels)
110	        self.qkv = nn.Conv1d(channels, channels * 3, 1)

151	            pad = 2
152	            self.conv = nn.Conv1d(self.channels, self.out_channels, ksize, padding=pad)
153	
154	    def forward(self, x):
155	        assert x.shape[1] == self.channels
156	        x = F.interpolate(x, scale_factor=self.factor, mode="nearest")
157	        if self.use_conv:
158	            x = self.conv(x)

176	        stride = factor
177	        if use_conv:
178	            self.op = nn.Conv1d(self.channels, self.out_channels, ksize, stride=stride, padding=pad)
179	        else:
180	            assert self.channels == self.out_channels
181	            self.op = nn.AvgPool1d(kernel_size=stride, stride=stride)
182	
183	    def forward(self, x):

180	            assert self.channels == self.out_channels
181	            self.op = nn.AvgPool1d(kernel_size=stride, stride=stride)
182	
183	    def forward(self, x):
184	        assert x.shape[1] == self.channels
185	        return self.op(x)
186	
187	

341	        if (
342	            len(inp.shape) == 3
343	        ):  # Automatically squeeze out the channels dimension if it is present (assuming mono-audio)
344	            inp = inp.squeeze(1)
345	        assert len(inp.shape) == 2
346	        self.mel_stft = self.mel_stft.to(inp.device)
347	        mel = self.mel_stft(inp)
348	        # Perform dynamic range compression

364	        self.wrap = wrap
365	
366	    def forward(self, x, *args, **kwargs):
367	        for k, v in kwargs.items():
368	            assert not (isinstance(v, torch.Tensor) and v.requires_grad)  # This would screw up checkpointing.
369	        partial = functools.partial(self.wrap, **kwargs)
370	        return partial(x, *args)
371	

38	    elif audiopath[-4:] == ".mp3":
39	        audio, lsr = librosa.load(audiopath, sr=None)
40	        audio = torch.FloatTensor(audio)
41	    else:
42	        assert False, f"Unsupported audio format provided: {audiopath[-4:]}"
43	
44	    # Remove any channel data.
45	    if len(audio.shape) > 1:

45	    if len(audio.shape) > 1:
46	        if audio.shape[0] < 5:
47	            audio = audio[0]
48	        else:
49	            assert audio.shape[1] < 5
50	            audio = audio[:, 0]
51	
52	    return audio, lsr

139	                print("Cannot combine a random voice with a non-random voice. Just using a random voice.")
140	            return None, None
141	        clip, latent = load_voice(voice, extra_voice_dirs)
142	        if latent is None:
143	            assert (
144	                len(latents) == 0
145	            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
146	            clips.extend(clip)
147	        elif clip is None:
148	            assert (

144	                len(latents) == 0
145	            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
146	            clips.extend(clip)
147	        elif clip is None:
148	            assert (
149	                len(clips) == 0
150	            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
151	            latents.append(latent)
152	    if len(latents) == 0:
153	        return clips, None

95	        output_attentions=None,
96	        output_hidden_states=None,
97	        return_dict=None,
98	    ):
99	        assert self.cached_mel_emb is not None
100	        assert inputs_embeds is None  # Not supported by this inference model.
101	        assert labels is None  # Training not supported by this inference model.
102	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

96	        output_hidden_states=None,
97	        return_dict=None,
98	    ):
99	        assert self.cached_mel_emb is not None
100	        assert inputs_embeds is None  # Not supported by this inference model.
101	        assert labels is None  # Training not supported by this inference model.
102	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
103	

97	        return_dict=None,
98	    ):
99	        assert self.cached_mel_emb is not None
100	        assert inputs_embeds is None  # Not supported by this inference model.
101	        assert labels is None  # Training not supported by this inference model.
102	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
103	
104	        # Create embedding

582	        trunc_index = fake_inputs.shape[1]
583	        if input_tokens is None:
584	            inputs = fake_inputs
585	        else:
586	            assert (
587	                num_return_sequences % input_tokens.shape[0] == 0
588	            ), "The number of return sequences must be divisible by the number of input sequences"
589	            fake_inputs = fake_inputs.repeat(num_return_sequences, 1)
590	            input_tokens = input_tokens.repeat(num_return_sequences // input_tokens.shape[0], 1)
591	            inputs = torch.cat([fake_inputs, input_tokens], dim=1)

39	    for obj in (mean1, logvar1, mean2, logvar2):
40	        if isinstance(obj, th.Tensor):
41	            tensor = obj
42	            break
43	    assert tensor is not None, "at least one argument must be a Tensor"
44	
45	    # Force variances to be Tensors. Broadcasting helps convert scalars to
46	    # Tensors, but it does not work for th.exp().

67	    :param means: the Gaussian mean Tensor.
68	    :param log_scales: the Gaussian log stddev Tensor.
69	    :return: a tensor like x of log probabilities (in nats).
70	    """
71	    assert x.shape == means.shape == log_scales.shape
72	    centered_x = x - means
73	    inv_stdv = th.exp(-log_scales)
74	    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)

82	        x < -0.999,
83	        log_cdf_plus,
84	        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
85	    )
86	    assert log_probs.shape == x.shape
87	    return log_probs
88	
89	

214	
215	        # Use float64 for accuracy.
216	        betas = np.array(betas, dtype=np.float64)
217	        self.betas = betas
218	        assert len(betas.shape) == 1, "betas must be 1-D"
219	        assert (betas > 0).all() and (betas <= 1).all()
220	
221	        self.num_timesteps = int(betas.shape[0])

215	        # Use float64 for accuracy.
216	        betas = np.array(betas, dtype=np.float64)
217	        self.betas = betas
218	        assert len(betas.shape) == 1, "betas must be 1-D"
219	        assert (betas > 0).all() and (betas <= 1).all()
220	
221	        self.num_timesteps = int(betas.shape[0])
222	

223	        alphas = 1.0 - betas
224	        self.alphas_cumprod = np.cumprod(alphas, axis=0)
225	        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
226	        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
227	        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
228	
229	        # calculations for diffusion q(x_t | x_{t-1}) and others
230	        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)

266	        :return: A noisy version of x_start.
267	        """
268	        if noise is None:
269	            noise = th.randn_like(x_start)
270	        assert noise.shape == x_start.shape
271	        return (
272	            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
273	            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise

279	
280	            q(x_{t-1} | x_t, x_0)
281	
282	        """
283	        assert x_start.shape == x_t.shape
284	        posterior_mean = (
285	            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
286	            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t

286	            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
287	        )
288	        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
289	        posterior_log_variance_clipped = _extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
290	        assert (
291	            posterior_mean.shape[0]
292	            == posterior_variance.shape[0]
293	            == posterior_log_variance_clipped.shape[0]
294	            == x_start.shape[0]
295	        )
296	        return posterior_mean, posterior_variance, posterior_log_variance_clipped
297	
298	    def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None):

319	        if model_kwargs is None:
320	            model_kwargs = {}
321	
322	        B, C = x.shape[:2]
323	        assert t.shape == (B,)
324	        model_output = model(x, self._scale_timesteps(t), **model_kwargs)
325	        if self.conditioning_free:
326	            model_output_no_conditioning = model(x, self._scale_timesteps(t), conditioning_free=True, **model_kwargs)

325	        if self.conditioning_free:
326	            model_output_no_conditioning = model(x, self._scale_timesteps(t), conditioning_free=True, **model_kwargs)
327	
328	        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
329	            assert model_output.shape == (B, C * 2, *x.shape[2:])
330	            model_output, model_var_values = th.split(model_output, C, dim=1)
331	            if self.conditioning_free:
332	                model_output_no_conditioning, _ = th.split(model_output_no_conditioning, C, dim=1)

357	            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
358	
359	        if self.conditioning_free:
360	            if self.ramp_conditioning_free:
361	                assert t.shape[0] == 1  # This should only be used in inference.
362	                cfk = self.conditioning_free_k * (1 - self._scale_timesteps(t)[0].item() / self.num_timesteps)
363	            else:
364	                cfk = self.conditioning_free_k

382	            model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
383	        else:
384	            raise NotImplementedError(self.model_mean_type)
385	
386	        assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
387	        return {
388	            "mean": model_mean,
389	            "variance": model_variance,

391	            "pred_xstart": pred_xstart,
392	        }
393	
394	    def _predict_xstart_from_eps(self, x_t, t, eps):
395	        assert x_t.shape == eps.shape
396	        return (
397	            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
398	            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps

398	            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
399	        )
400	
401	    def _predict_xstart_from_xprev(self, x_t, t, xprev):
402	        assert x_t.shape == xprev.shape
403	        return (  # (xprev - coef2*x_t) / coef1
404	            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
405	            - _extract_into_tensor(self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape) * x_t

461	        device=None,  # ALL UNUSED
462	        model_kwargs=None,  # {'precomputed_aligned_embeddings': precomputed_embeddings},
463	        progress=False,  # unused as well
464	    ):
465	        assert isinstance(model_kwargs, dict)
466	        if device is None:
467	            device = next(model.parameters()).device
468	        s_in = noise.new_ones([noise.shape[0]])

652	        p_sample().
653	        """
654	        if device is None:
655	            device = next(model.parameters()).device
656	        assert isinstance(shape, (tuple, list))
657	        if noise is not None:
658	            img = noise
659	        else:

728	    ):
729	        """
730	        Sample x_{t+1} from the model using DDIM reverse ODE.
731	        """
732	        assert eta == 0.0, "Reverse ODE only for deterministic path"
733	        out = self.p_mean_variance(
734	            model,
735	            x,

804	        Same usage as p_sample_loop_progressive().
805	        """
806	        if device is None:
807	            device = next(model.parameters()).device
808	        assert isinstance(shape, (tuple, list))
809	        if noise is not None:
810	            img = noise
811	        else:

852	
853	        decoder_nll = -discretized_gaussian_log_likelihood(
854	            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
855	        )
856	        assert decoder_nll.shape == x_start.shape
857	        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
858	
859	        # At the first timestep return the decoder NLL,

906	                ModelVarType.LEARNED,
907	                ModelVarType.LEARNED_RANGE,
908	            ]:
909	                B, C = x_t.shape[:2]
910	                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
911	                model_output, model_var_values = th.split(model_output, C, dim=1)
912	                # Learn the variance using the variational bound, but don't let
913	                # it affect our mean prediction.

934	                target = noise
935	                x_start_pred = self._predict_xstart_from_eps(x_t, t, model_output)
936	            else:
937	                raise NotImplementedError(self.model_mean_type)
938	            assert model_output.shape == target.shape == x_start.shape
939	            terms["mse"] = mean_flat((target - model_output) ** 2)
940	            terms["x_start_predicted"] = x_start_pred
941	            if "vb" in terms:

968	            noise = th.randn_like(x_start)
969	        x_t = self.q_sample(x_start, t, noise=noise)
970	        terms = {}
971	        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
972	            assert False  # not currently supported for this type of diffusion.
973	        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
974	            model_outputs = model(x_t, x_start, self._scale_timesteps(t), **model_kwargs)
975	            terms.update({k: o for k, o in zip(model_output_keys, model_outputs)})

978	                ModelVarType.LEARNED,
979	                ModelVarType.LEARNED_RANGE,
980	            ]:
981	                B, C = x_t.shape[:2]
982	                assert model_output.shape == (B, C, 2, *x_t.shape[2:])
983	                model_output, model_var_values = model_output[:, :, 0], model_output[:, :, 1]
984	                # Learn the variance using the variational bound, but don't let
985	                # it affect our mean prediction.

1006	                target = noise
1007	                x_start_pred = self._predict_xstart_from_eps(x_t, t, model_output)
1008	            else:
1009	                raise NotImplementedError(self.model_mean_type)
1010	            assert model_output.shape == target.shape == x_start.shape
1011	            terms["mse"] = mean_flat((target - model_output) ** 2)
1012	            terms["x_start_predicted"] = x_start_pred
1013	            if "vb" in terms:

339	        :param precomputed_aligned_embeddings: Embeddings returned from self.timestep_independent()
340	        :param conditioning_free: When set, all conditioning inputs (including tokens and conditioning_input) will not be considered.
341	        :return: an [N x C x ...] Tensor of outputs.
342	        """
343	        assert precomputed_aligned_embeddings is not None or (
344	            aligned_conditioning is not None and conditioning_latent is not None
345	        )
346	        assert not (
347	            return_code_pred and precomputed_aligned_embeddings is not None
348	        )  # These two are mutually exclusive.

342	        """
343	        assert precomputed_aligned_embeddings is not None or (
344	            aligned_conditioning is not None and conditioning_latent is not None
345	        )
346	        assert not (
347	            return_code_pred and precomputed_aligned_embeddings is not None
348	        )  # These two are mutually exclusive.
349	
350	        unused_params = []
351	        if conditioning_free:

379	                self.training
380	                and self.layer_drop > 0
381	                and i != 0
382	                and i != (len(self.layers) - 1)
383	                and random.random() < self.layer_drop
384	            ):
385	                unused_params.extend(list(lyr.parameters()))
386	            else:

104	        if schedule == "discrete":
105	            if betas is not None:
106	                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
107	            else:
108	                assert alphas_cumprod is not None
109	                log_alphas = 0.5 * torch.log(alphas_cumprod)
110	            self.total_N = len(log_alphas)
111	            self.T = 1.0

355	        """
356	        if guidance_type == "uncond":
357	            return noise_pred_fn(x, t_continuous)
358	        elif guidance_type == "classifier":
359	            assert classifier_fn is not None
360	            t_input = get_model_input_time(t_continuous)
361	            cond_grad = cond_grad_fn(x, t_input)
362	            sigma_t = noise_schedule.marginal_std(t_continuous)

371	                c_in = torch.cat([unconditional_condition, condition])
372	                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
373	                return noise_uncond + guidance_scale * (noise - noise_uncond)
374	
375	    assert model_type in ["noise", "x_start", "v", "score"]
376	    assert guidance_type in ["uncond", "classifier", "classifier-free"]
377	    return model_fn
378	

372	                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
373	                return noise_uncond + guidance_scale * (noise - noise_uncond)
374	
375	    assert model_type in ["noise", "x_start", "v", "score"]
376	    assert guidance_type in ["uncond", "classifier", "classifier-free"]
377	    return model_fn
378	
379	

445	            with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
446	        """
447	        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
448	        self.noise_schedule = noise_schedule
449	        assert algorithm_type in ["dpmsolver", "dpmsolver++"]
450	        self.algorithm_type = algorithm_type
451	        if correcting_x0_fn == "dynamic_thresholding":
452	            self.correcting_x0_fn = self.dynamic_thresholding_fn

1215	        For discrete-time DPMs, we use `t_start=1/N`, where `N` is the total time steps during training.
1216	        """
1217	        t_0 = 1.0 / self.noise_schedule.total_N if t_start is None else t_start
1218	        t_T = self.noise_schedule.T if t_end is None else t_end
1219	        assert (
1220	            t_0 > 0 and t_T > 0
1221	        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
1222	        return self.sample(
1223	            x,
1224	            steps=steps,

1360	
1361	        """
1362	        t_0 = 1.0 / self.noise_schedule.total_N if t_end is None else t_end
1363	        t_T = self.noise_schedule.T if t_start is None else t_start
1364	        assert (
1365	            t_0 > 0 and t_T > 0
1366	        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
1367	        if return_intermediate:
1368	            assert method in [
1369	                "multistep",

1364	        assert (
1365	            t_0 > 0 and t_T > 0
1366	        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
1367	        if return_intermediate:
1368	            assert method in [
1369	                "multistep",
1370	                "singlestep",
1371	                "singlestep_fixed",
1372	            ], "Cannot use adaptive solver when saving intermediate values"
1373	        if self.correcting_xt_fn is not None:
1374	            assert method in [
1375	                "multistep",

1370	                "singlestep",
1371	                "singlestep_fixed",
1372	            ], "Cannot use adaptive solver when saving intermediate values"
1373	        if self.correcting_xt_fn is not None:
1374	            assert method in [
1375	                "multistep",
1376	                "singlestep",
1377	                "singlestep_fixed",
1378	            ], "Cannot use adaptive solver when correcting_xt_fn is not None"
1379	        device = x.device
1380	        intermediates = []
1381	        with torch.no_grad():

1389	                    rtol=rtol,
1390	                    solver_type=solver_type,
1391	                )
1392	            elif method == "multistep":
1393	                assert steps >= order
1394	                timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1395	                assert timesteps.shape[0] - 1 == steps
1396	                # Init the initial values.

1391	                )
1392	            elif method == "multistep":
1393	                assert steps >= order
1394	                timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
1395	                assert timesteps.shape[0] - 1 == steps
1396	                # Init the initial values.
1397	                step = 0
1398	                t = timesteps[step]

45	# classes
46	class SequentialSequence(nn.Module):
47	    def __init__(self, layers, args_route={}, layer_dropout=0.0):
48	        super().__init__()
49	        assert all(
50	            len(route) == len(layers) for route in args_route.values()
51	        ), "each argument route map must have the same depth as the number of sequential layers"
52	        self.layers = layers
53	        self.args_route = args_route
54	        self.layer_dropout = layer_dropout

243	            (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length).
244	        """
245	        batch, _, in_length = x.shape
246	        batch, _, out_channels, kernel_size, kernel_length = kernel.shape
247	        assert in_length == (kernel_length * hop_size), "length of (x, kernel) is not matched"
248	
249	        padding = dilation * int((kernel_size - 1) / 2)
250	        x = F.pad(x, (padding, padding), "constant", 0)  # (batch, in_channels, in_length + 2*padding)

398	    print(c.shape)
399	
400	    y = model(c, z)
401	    print(y.shape)
402	    assert y.shape == torch.Size([3, 1, 2560])
403	
404	    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
405	    print(pytorch_total_params)

11	    Finally got to use my DP skills!
12	    """
13	    if record is None:
14	        record = {}
15	    assert skip_character not in s1, f"Found the skip character {skip_character} in the provided string, {s1}"
16	    if len(s1) == 0:
17	        return ""
18	    if len(s2) == 0:

102	
103	        pop_till_you_win()
104	        if not (len(expected_tokens) == 0 and len(alignments) == len(expected_text)):
105	            torch.save([audio, expected_text], "alignment_debug.pth")
106	            assert False, (
107	                "Something went wrong with the alignment algorithm. I've dumped a file, 'alignment_debug.pth' to"
108	                "your current working directory. Please report this along with the file so it can get fixed."
109	            )
110	
111	        # Now fix up alignments. Anything with -1 should be interpolated.
112	        alignments.append(orig_len)  # This'll get removed but makes the algorithm below more readable.

127	            return audio
128	        splitted = expected_text.split("[")
129	        fully_split = [splitted[0]]
130	        for spl in splitted[1:]:
131	            assert "]" in spl, 'Every "[" character must be paired with a "]" with no nesting.'
132	            fully_split.extend(spl.split("]"))
133	
134	        # At this point, fully_split is a list of strings, with every other string being something that should be redacted.

575	        self.to_out = nn.Sequential(nn.Linear(v_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(v_dim, dim)
576	
577	        self.rel_pos_bias = rel_pos_bias
578	        if rel_pos_bias:
579	            assert (
580	                rel_pos_num_buckets <= rel_pos_max_distance
581	            ), "number of relative position buckets must be less than the relative position max distance"
582	            self.rel_pos = RelativePositionBias(
583	                scale=dim_head**0.5,
584	                causal=causal,

695	            dots.masked_fill_(~input_mask, mask_value)
696	            del input_mask
697	
698	        if exists(attn_mask):
699	            assert (
700	                2 <= attn_mask.ndim <= 4
701	            ), "attention mask must have greater than 2 dimensions but less than or equal to 4"
702	            if attn_mask.ndim == 2:
703	                attn_mask = rearrange(attn_mask, "i j -> () () i j")
704	            elif attn_mask.ndim == 3:

805	
806	        rotary_emb_dim = max(default(rotary_emb_dim, dim_head // 2), 32)
807	        self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim) if rotary_pos_emb else None
808	
809	        assert not (
810	            alibi_pos_bias and rel_pos_bias
811	        ), "you can only choose Alibi positional bias or T5 relative positional bias, not both"
812	
813	        if alibi_pos_bias:
814	            alibi_num_heads = default(alibi_num_heads, heads)

811	        ), "you can only choose Alibi positional bias or T5 relative positional bias, not both"
812	
813	        if alibi_pos_bias:
814	            alibi_num_heads = default(alibi_num_heads, heads)
815	            assert alibi_num_heads <= heads, "number of ALiBi heads must be less than the total number of heads"
816	            alibi_pos_klass = LearnedAlibiPositionalBias if alibi_learned or not causal else AlibiPositionalBias
817	            self.rel_pos = alibi_pos_klass(heads=alibi_num_heads, bidirectional=not causal)
818	        else:

817	            self.rel_pos = alibi_pos_klass(heads=alibi_num_heads, bidirectional=not causal)
818	        else:
819	            self.rel_pos = None
820	
821	        assert not (not pre_norm and sandwich_norm), "sandwich norm cannot be used when not using prenorm"
822	        self.pre_norm = pre_norm
823	        self.sandwich_norm = sandwich_norm
824	

865	        if exists(custom_layers):
866	            layer_types = custom_layers
867	        elif exists(par_ratio):
868	            par_depth = depth * len(default_block)
869	            assert 1 < par_ratio <= par_depth, "par ratio out of range"
870	            default_block = tuple(filter(not_equals("f"), default_block))
871	            par_attn = par_depth // par_ratio
872	            depth_cut = par_depth * 2 // 3  # 2 / 3 attention layer cutoff suggested by PAR paper

870	            default_block = tuple(filter(not_equals("f"), default_block))
871	            par_attn = par_depth // par_ratio
872	            depth_cut = par_depth * 2 // 3  # 2 / 3 attention layer cutoff suggested by PAR paper
873	            par_width = (depth_cut + depth_cut // par_attn) // par_attn
874	            assert len(default_block) <= par_width, "default block is too large for par_ratio"
875	            par_block = default_block + ("f",) * (par_width - len(default_block))
876	            par_head = par_block * par_attn
877	            layer_types = par_head + ("f",) * (par_depth - len(par_head))

875	            par_block = default_block + ("f",) * (par_width - len(default_block))
876	            par_head = par_block * par_attn
877	            layer_types = par_head + ("f",) * (par_depth - len(par_head))
878	        elif exists(sandwich_coef):
879	            assert sandwich_coef > 0 and sandwich_coef <= depth, "sandwich coefficient should be less than the depth"
880	            layer_types = ("a",) * sandwich_coef + default_block * (depth - sandwich_coef) + ("f",) * sandwich_coef
881	        else:
882	            layer_types = default_block * depth

937	        norm_scale_shift_inp=None,
938	        past_key_values=None,
939	        expected_seq_len=None,
940	    ):
941	        assert not (
942	            self.cross_attend ^ (exists(context) or exists(full_context))
943	        ), "context must be passed in if cross_attend is set to True"
944	        assert context is None or full_context is None, "only one of full_context or context can be provided"
945	
946	        hiddens = []

940	    ):
941	        assert not (
942	            self.cross_attend ^ (exists(context) or exists(full_context))
943	        ), "context must be passed in if cross_attend is set to True"
944	        assert context is None or full_context is None, "only one of full_context or context can be provided"
945	
946	        hiddens = []
947	        intermediates = []

955	
956	        rotary_pos_emb = None
957	        if exists(self.rotary_pos_emb):
958	            if not self.training and self.causal:
959	                assert (
960	                    expected_seq_len is not None
961	                ), "To decode a transformer with rotary embeddings, you must specify an `expected_seq_len`"
962	            elif expected_seq_len is None:
963	                expected_seq_len = 0
964	            seq_len = x.shape[1]

1050	
1051	
1052	class Encoder(AttentionLayers):
1053	    def __init__(self, **kwargs):
1054	        assert "causal" not in kwargs, "cannot set causality on encoder"
1055	        super().__init__(causal=False, **kwargs)
1056	
1057	

1056	
1057	
1058	class Decoder(AttentionLayers):
1059	    def __init__(self, **kwargs):
1060	        assert "causal" not in kwargs, "cannot set causality on decoder"
1061	        super().__init__(causal=True, **kwargs)
1062	
1063	

1068	
1069	class ViTransformerWrapper(nn.Module):
1070	    def __init__(self, *, image_size, patch_size, attn_layers, num_classes=None, dropout=0.0, emb_dropout=0.0):
1071	        super().__init__()
1072	        assert isinstance(attn_layers, Encoder), "attention layers must be an Encoder"
1073	        assert image_size % patch_size == 0, "image dimensions must be divisible by the patch size"
1074	        dim = attn_layers.dim
1075	        num_patches = (image_size // patch_size) ** 2

1069	class ViTransformerWrapper(nn.Module):
1070	    def __init__(self, *, image_size, patch_size, attn_layers, num_classes=None, dropout=0.0, emb_dropout=0.0):
1071	        super().__init__()
1072	        assert isinstance(attn_layers, Encoder), "attention layers must be an Encoder"
1073	        assert image_size % patch_size == 0, "image dimensions must be divisible by the patch size"
1074	        dim = attn_layers.dim
1075	        num_patches = (image_size // patch_size) ** 2
1076	        patch_dim = 3 * patch_size**2

1122	        tie_embedding=False,
1123	        use_pos_emb=True,
1124	    ):
1125	        super().__init__()
1126	        assert isinstance(attn_layers, AttentionLayers), "attention layers must be one of Encoder or Decoder"
1127	
1128	        dim = attn_layers.dim
1129	        emb_dim = default(emb_dim, dim)

1214	    def __init__(
1215	        self, *, max_seq_len, attn_layers, dim_in=None, dim_out=None, emb_dim=None, emb_dropout=0.0, use_pos_emb=True
1216	    ):
1217	        super().__init__()
1218	        assert isinstance(attn_layers, AttentionLayers), "attention layers must be one of Encoder or Decoder"
1219	
1220	        dim = attn_layers.dim
1221	

82	        Shapes:
83	            - x: :math:`[B, T]`
84	            - x_length: :math:`[B]`
85	        """
86	        assert x.shape[0] == x_lengths.shape[0]
87	        x = self.emb(x) * math.sqrt(self.hidden_channels)  # [b, t, h]
88	
89	        # concat the lang emb in embedding chars

111	        dropout_p=0,
112	        cond_channels=0,
113	        mean_only=False,
114	    ):
115	        assert channels % 2 == 0, "channels should be divisible by 2"
116	        super().__init__()
117	        self.half_channels = channels // 2
118	        self.mean_only = mean_only

239	        x = self.proj(x) * x_mask
240	
241	        if not reverse:
242	            flows = self.flows
243	            assert dr is not None
244	
245	            # condition encoder duration
246	            h = self.post_pre(dr)

164	        )
165	        c = -input_delta * (inputs - input_cumheights)
166	
167	        discriminant = b.pow(2) - 4 * a * c
168	        assert (discriminant >= 0).all()
169	
170	        root = (2 * c) / (-b - torch.sqrt(discriminant))
171	        outputs = root * input_bin_widths + input_cumwidths

185	
186	class UpsampledConv(nn.Module):
187	    def __init__(self, conv, *args, **kwargs):
188	        super().__init__()
189	        assert "stride" in kwargs.keys()
190	        self.stride = kwargs["stride"]
191	        del kwargs["stride"]
192	        self.conv = conv(*args, **kwargs)

230	        self.discrete_loss = DiscretizationLoss(
231	            num_tokens, 2, 1 / (num_tokens * 2), discretization_loss_averaging_steps
232	        )
233	
234	        assert positional_dims > 0 and positional_dims < 3  # This VAE only supports 1d and 2d inputs for now.
235	        if positional_dims == 2:
236	            conv = nn.Conv2d
237	            conv_transpose = nn.ConvTranspose2d

245	            act = nn.ReLU
246	        elif activation == "silu":
247	            act = nn.SiLU
248	        else:
249	            assert NotImplementedError()
250	
251	        enc_layers = []
252	        dec_layers = []

30	
31	    def forward(self, x):
32	        sl = x.shape[1]
33	        if self.relative:
34	            start = random.randint(sl, self.seq_len) - sl
35	            return self.emb(torch.arange(start, start + sl, device=x.device))
36	        else:
37	            return self.emb(torch.arange(0, sl, device=x.device))

341	                for i in range(prompt_codes.shape[0]):
342	                    if lengths[i] < prompt_len:
343	                        start = 0
344	                    else:
345	                        start = random.randint(0, lengths[i] - prompt_len)
346	                prompt = prompt_codes[:, start : start + prompt_len]
347	
348	        # add start and stop tokens

395	        If return_latent is specified, loss & logits are not computed or returned. Only the predicted latents are returned.
396	        """
397	        # ❗ FIXIT
398	        if self.max_conditioning_inputs == 0:
399	            assert cond_mels is None, " ❗ cond_mels is not None, but max_conditioning_inputs == 0"
400	
401	        max_text_len = text_lengths.max()
402	        code_lengths = torch.ceil(wav_lengths / self.code_stride_len).long() + 3

426	        if max_mel_len > audio_codes.shape[-1]:
427	            audio_codes = F.pad(audio_codes, (0, max_mel_len - audio_codes.shape[-1]))
428	
429	        # 💖 Lovely assertions
430	        assert (
431	            max_mel_len <= audio_codes.shape[-1]
432	        ), f" ❗ max_mel_len ({max_mel_len}) > audio_codes.shape[-1] ({audio_codes.shape[-1]})"
433	        assert (
434	            max_text_len <= text_inputs.shape[-1]
435	        ), f" ❗ max_text_len ({max_text_len}) > text_inputs.shape[-1] ({text_inputs.shape[-1]})"

429	        # 💖 Lovely assertions
430	        assert (
431	            max_mel_len <= audio_codes.shape[-1]
432	        ), f" ❗ max_mel_len ({max_mel_len}) > audio_codes.shape[-1] ({audio_codes.shape[-1]})"
433	        assert (
434	            max_text_len <= text_inputs.shape[-1]
435	        ), f" ❗ max_text_len ({max_text_len}) > text_inputs.shape[-1] ({text_inputs.shape[-1]})"
436	
437	        # Append stop token to text inputs
438	        text_inputs = F.pad(text_inputs[:, :max_text_len], (0, 1), value=self.stop_text_token)

533	        for idx, l in enumerate(code_lengths):
534	            mel_targets[idx, l + 1 :] = -1
535	
536	        # check if stoptoken is in every row of mel_targets
537	        assert (mel_targets == self.stop_audio_token).sum() >= mel_targets.shape[
538	            0
539	        ], f" ❗ mel_targets does not contain stop token ({self.stop_audio_token}) in every row."
540	
541	        # ignore the loss for the segment used for conditioning
542	        # coin flip for the segment to be ignored

68	        output_attentions=None,
69	        output_hidden_states=None,
70	        return_dict=None,
71	    ):
72	        assert self.cached_prefix_emb is not None
73	        assert inputs_embeds is None  # Not supported by this inference model.
74	        assert labels is None  # Training not supported by this inference model.
75	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

69	        output_hidden_states=None,
70	        return_dict=None,
71	    ):
72	        assert self.cached_prefix_emb is not None
73	        assert inputs_embeds is None  # Not supported by this inference model.
74	        assert labels is None  # Training not supported by this inference model.
75	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
76	

70	        return_dict=None,
71	    ):
72	        assert self.cached_prefix_emb is not None
73	        assert inputs_embeds is None  # Not supported by this inference model.
74	        assert labels is None  # Training not supported by this inference model.
75	        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
76	
77	        # assert len(past_key_values) + len(input_ids) == attention_mask.shape[1]

330	        state = torch.load(checkpoint_path, map_location=torch.device("cpu"))
331	        self.load_state_dict(state["model"])
332	        if eval:
333	            self.eval()
334	            assert not self.training
335	            self.remove_weight_norm()
336	
337	

411	        self.coefficient = coefficient
412	        self.register_buffer("filter", torch.FloatTensor([-self.coefficient, 1.0]).unsqueeze(0).unsqueeze(0))
413	
414	    def forward(self, x):
415	        assert len(x.size()) == 2
416	
417	        x = torch.nn.functional.pad(x.unsqueeze(1), (1, 0), "reflect")
418	        return torch.nn.functional.conv1d(x, self.filter).squeeze(1)

604	                criterion = criterion.cuda()
605	
606	        if eval:
607	            self.eval()
608	            assert not self.training
609	
610	        if not eval:
611	            return criterion, state["step"]

727	
728	        self.load_state_dict(state)
729	        if eval:
730	            self.eval()
731	            assert not self.training
732	            self.waveform_decoder.remove_weight_norm()

29	    elif channels <= 64:
30	        groups = 16
31	    while channels % groups != 0:
32	        groups = int(groups / 2)
33	    assert groups > 2
34	    return GroupNorm32(groups, channels)
35	
36	

52	        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
53	        :return: an [N x (H * C) x T] tensor after attention.
54	        """
55	        bs, width, length = qkv.shape
56	        assert width % (3 * self.n_heads) == 0
57	        ch = width // (3 * self.n_heads)
58	        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
59	        scale = 1 / math.sqrt(math.sqrt(ch))

85	        self.do_activation = do_activation
86	        if num_head_channels == -1:
87	            self.num_heads = num_heads
88	        else:
89	            assert (
90	                channels % num_head_channels == 0
91	            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
92	            self.num_heads = channels // num_head_channels
93	        self.norm = normalization(channels)
94	        self.qkv = conv_nd(1, channels, out_channels * 3, 1)

43	        self.causal = causal
44	        self.register_buffer("mask", None, persistent=False)
45	
46	        self.use_flash = use_flash
47	        assert not (
48	            use_flash and version.parse(torch.__version__) < version.parse("2.0.0")
49	        ), "in order to use flash attention, you must be using pytorch 2.0 or above"
50	
51	        # determine efficient attention configs for cuda and cpu
52	        self.config = namedtuple("EfficientAttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"])

180	
181	        if not self.cond:
182	            return out
183	
184	        assert exists(cond)
185	        gamma, beta = self.to_gamma_beta(cond).chunk(2, dim=-1)
186	        gamma, beta = map(lambda t: rearrange(t, "b d -> b 1 d"), (gamma, beta))
187	        return out * gamma + beta

193	        (kernel_size,) = self.kernel_size
194	        (dilation,) = self.dilation
195	        (stride,) = self.stride
196	
197	        assert stride == 1
198	        self.causal_padding = dilation * (kernel_size - 1)
199	
200	    def forward(self, x):

530	        try:
531	            text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
532	            text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
533	            text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
534	        except:
535	            pass
536	        if lang != "tr":
537	            text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
538	        text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)

763	        ("이것은 1 번째 테스트입니다", "이것은 첫 번째 테스트입니다", "ko"),
764	    ]
765	    for a, b, lang in test_cases:
766	        out = expand_numbers_multilingual(a, lang=lang)
767	        assert out == b, f"'{out}' vs '{b}'"
768	
769	
770	def test_abbreviations_multilingual():

806	    ]
807	
808	    for a, b, lang in test_cases:
809	        out = expand_abbreviations_multilingual(a, lang=lang)
810	        assert out == b, f"'{out}' vs '{b}'"
811	
812	
813	def test_symbols_multilingual():

833	    ]
834	
835	    for a, b, lang in test_cases:
836	        out = expand_symbols_multilingual(a, lang=lang)
837	        assert out == b, f"'{out}' vs '{b}'"
838	
839	
840	if __name__ == "__main__":

15	    """Returns a dictionary of samples keyed by language."""
16	    samples_by_col = {}
17	    for sample in samples:
18	        col_val = sample[col]
19	        assert isinstance(col_val, str)
20	        if col_val not in samples_by_col:
21	            samples_by_col[col_val] = []
22	        samples_by_col[col_val].append(sample)

28	    # if eval uses a middle size sample when it is possible to be more reproducible
29	    if is_eval:
30	        sample_length = int((min_sample_length + max_sample_length) / 2)
31	    else:
32	        sample_length = random.randint(min_sample_length, max_sample_length)
33	    gap = rel_clip.shape[-1] - sample_length
34	    if gap < 0:
35	        sample_length = rel_clip.shape[-1] // 2

38	    # if eval start always from the position 0 to be more reproducible
39	    if is_eval:
40	        rand_start = 0
41	    else:
42	        rand_start = random.randint(0, gap)
43	
44	    rand_end = rand_start + sample_length
45	    rel_clip = rel_clip[:, rand_start:rand_end]

61	        self.sample_rate = sample_rate
62	        self.max_wav_len = model_args.max_wav_length
63	        self.max_text_len = model_args.max_text_length
64	        self.use_masking_gt_prompt_approach = model_args.gpt_use_masking_gt_prompt_approach
65	        assert self.max_wav_len is not None and self.max_text_len is not None
66	
67	        self.samples = samples
68	        if not is_eval:

81	        new_samples = []
82	        for sample in self.samples:
83	            try:
84	                tseq, _, wav, _, _, _ = self.load_item(sample)
85	            except:
86	                continue
87	            # Basically, this audio file is nonexistent or too long to be supported by the dataset.
88	            if (
89	                wav is None

97	
98	    def get_text(self, text, lang):
99	        tokens = self.tokenizer.encode(text, lang)
100	        tokens = torch.IntTensor(tokens)
101	        assert not torch.any(tokens == 1), f"UNK token found in {text} -> {self.tokenizer.decode(tokens)}"
102	        # The stop token should always be sacred.
103	        assert not torch.any(tokens == 0), f"Stop token found in {text}"
104	        return tokens

99	        tokens = self.tokenizer.encode(text, lang)
100	        tokens = torch.IntTensor(tokens)
101	        assert not torch.any(tokens == 1), f"UNK token found in {text} -> {self.tokenizer.decode(tokens)}"
102	        # The stop token should always be sacred.
103	        assert not torch.any(tokens == 0), f"Stop token found in {text}"
104	        return tokens
105	
106	    def load_item(self, sample):

140	            sample = self.samples[index]
141	            sample_id = str(index)
142	        else:
143	            # select a random language
144	            lang = random.choice(list(self.samples.keys()))
145	            # select random sample
146	            index = random.randint(0, len(self.samples[lang]) - 1)
147	            sample = self.samples[lang][index]

142	        else:
143	            # select a random language
144	            lang = random.choice(list(self.samples.keys()))
145	            # select random sample
146	            index = random.randint(0, len(self.samples[lang]) - 1)
147	            sample = self.samples[lang][index]
148	            # a unique id for each sampel to deal with fails
149	            sample_id = lang + "_" + str(index)

56	
57	def callback_clearml_load_save(operation_type, model_info):
58	    # return None means skip the file upload/log, returning model_info will continue with the log/upload
59	    # you can also change the upload destination file name model_info.upload_filename or check the local file size with Path(model_info.local_model_path).stat().st_size
60	    assert operation_type in ("load", "save")
61	    # print(operation_type, model_info.__dict__)
62	
63	    if "similarities.pth" in model_info.__dict__["local_model_path"]:

489	
490	        if eval:
491	            self.xtts.gpt.init_gpt_for_inference(kv_cache=self.args.kv_cache, use_deepspeed=False)
492	            self.eval()
493	            assert not self.training
494	
495	    @staticmethod
496	    def init_from_config(config: "GPTTrainerConfig", samples: Union[List[List], List[Dict]] = None):

1168	    ndone = 0
1169	    with open(args.ifile, "r", encoding="utf8") as istream, open(args.ofile, "w+", encoding="utf8") as ostream:
1170	        if args.format == "tsv":
1171	            reader = csv.DictReader(istream, delimiter="\t")
1172	            assert "TEXT" in reader.fieldnames
1173	            print("\t".join(reader.fieldnames), file=ostream)
1174	
1175	            for item in reader:

402	        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
403	        self.load_state_dict(state["model"])
404	        if eval:
405	            self.eval()
406	            assert not self.training
407	
408	    def get_criterion(self):
409	        from TTS.tts.layers.losses import AlignTTSLoss  # pylint: disable=import-outside-toplevel

116	            self.decoder.set_r(config.r)
117	        if eval:
118	            self.eval()
119	            print(f" > Model's reduction rate `r` is set to: {self.decoder.r}")
120	            assert not self.training
121	
122	    def get_criterion(self) -> nn.Module:
123	        """Get the model criterion used in training."""

205	                # by cutting off from the largest duration indeces.
206	                extra_frames = dur.sum() - mel_lengths[idx]
207	                largest_idxs = torch.argsort(-dur)[:extra_frames]
208	                dur[largest_idxs] -= 1
209	                assert (
210	                    dur.sum() == mel_lengths[idx]
211	                ), f" [!] total duration {dur.sum()} vs spectrogram length {mel_lengths[idx]}"
212	                durations[idx, : text_lengths[idx]] = dur
213	
214	        # set stop targets wrt reduction factor

95	    return torch.ceil(lens / stride).int()
96	
97	
98	def initialize_embeddings(shape: Tuple[int]) -> torch.Tensor:
99	    assert len(shape) == 2, "Can only initialize 2-D embedding matrices ..."
100	    return torch.randn(shape) * np.sqrt(2 / shape[1])
101	
102	

135	    """
136	    x, sr = torchaudio.load(
137	        file_path,
138	    )
139	    assert (x > 1).sum() + (x < -1).sum() == 0
140	    return x, sr
141	
142	

1546	        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
1547	        self.load_state_dict(state["model"])
1548	        if eval:
1549	            self.eval()
1550	            assert not self.training
1551	
1552	    def get_state_dict(self):
1553	        """Custom state dict of the model with all the necessary components for inference."""

833	        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
834	        self.load_state_dict(state["model"])
835	        if eval:
836	            self.eval()
837	            assert not self.training
838	
839	    def get_criterion(self):
840	        from TTS.tts.layers.losses import ForwardTTSLoss  # pylint: disable=import-outside-toplevel

121	            self.embedded_speaker_dim = (
122	                config.d_vector_dim if "d_vector_dim" in config and config.d_vector_dim is not None else 512
123	            )
124	            if self.speaker_manager is not None:
125	                assert (
126	                    config.d_vector_dim == self.speaker_manager.embedding_dim
127	                ), " [!] d-vector dimension mismatch b/w config and speaker manager."
128	        # init speaker embedding layer
129	        if config.use_speaker_embedding and not config.use_d_vector_file:
130	            print(" > Init speaker_embedding layer.")

526	        self.load_state_dict(state["model"])
527	        if eval:
528	            self.eval()
529	            self.store_inverse()
530	            assert not self.training
531	
532	    @staticmethod
533	    def get_criterion():

257	        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
258	        self.load_state_dict(state["model"])
259	        if eval:
260	            self.eval()
261	            assert not self.training
262	
263	    def on_init_start(self, trainer):
264	        """If the current dataset does not have normalisation statistics and initialisation transition_probability it computes them otherwise loads."""

273	        self.load_state_dict(state["model"])
274	        if eval:
275	            self.eval()
276	            self.decoder.store_inverse()
277	            assert not self.training
278	
279	    def on_init_start(self, trainer):
280	        """If the current dataset does not have normalisation statistics and initialisation transition_probability it computes them otherwise loads."""

79	    gap = clip.shape[-1] - cond_length
80	    if gap < 0:
81	        clip = F.pad(clip, pad=(0, abs(gap)))
82	    elif gap > 0:
83	        rand_start = random.randint(0, gap)
84	        clip = clip[:, rand_start : rand_start + cond_length]
85	    mel_clip = TorchMelSpectrogram(**kwargs)(clip.unsqueeze(0)).squeeze(0)
86	    return mel_clip.unsqueeze(0).to(device)

418	            0 - latents will be generated as in original tortoise, using ~4.27s from each voice sample, averaging latent across all samples
419	            1 - latents will be generated using (almost) entire voice samples, averaged across all the ~4.27s chunks
420	            2 - latents will be generated using (almost) entire voice samples, averaged per voice sample
421	        """
422	        assert latent_averaging_mode in [
423	            0,
424	            1,
425	            2,
426	        ], "latent_averaging mode has to be one of (0, 1, 2)"
427	
428	        with torch.no_grad():
429	            voice_samples = [[v.to(self.device) for v in ls] for ls in voice_samples]

670	        deterministic_seed = deterministic_state(seed=use_deterministic_seed)
671	
672	        text_tokens = torch.IntTensor(self.tokenizer.encode(text)).unsqueeze(0).to(self.device)
673	        text_tokens = F.pad(text_tokens, (0, 1))  # This may not be necessary.
674	        assert (
675	            text_tokens.shape[-1] < 400
676	        ), "Too much text provided. Break the text up into separate segments and re-try inference."
677	
678	        if voice_samples is not None:
679	            (

70	    Return Shapes:
71	        - x: :math:`[1, T]`
72	    """
73	    x, sr = torchaudio.load(file_path)
74	    assert (x > 1).sum() + (x < -1).sum() == 0
75	    return x, sr
76	
77	

241	    dataset_samples_weight = dataset_samples_weight / np.linalg.norm(dataset_samples_weight)
242	    if multi_dict is not None:
243	        # check if all keys are in the multi_dict
244	        for k in multi_dict:
245	            assert k in unique_attr_names, f"{k} not in {unique_attr_names}"
246	        # scale weights
247	        multiplier_samples = np.array([multi_dict.get(item[attr_name], 1.0) for item in items])
248	        dataset_samples_weight *= multiplier_samples

1138	                    x, x_mask, g=g if self.args.condition_dp_on_speaker else None, lang_emb=lang_emb
1139	                )
1140	            w = torch.exp(logw) * x_mask * self.length_scale
1141	        else:
1142	            assert durations.shape[-1] == x.shape[-1]
1143	            w = durations.unsqueeze(0)
1144	
1145	        w_ceil = torch.ceil(w)

1209	            y_lengths (Tensor): Length of each reference spectrogram. Tensor of shape [B]
1210	            speaker_cond_src (Tensor): Reference speaker ID. Tensor of shape [B,]
1211	            speaker_cond_tgt (Tensor): Target speaker ID. Tensor of shape [B,]
1212	        """
1213	        assert self.num_speakers > 0, "num_speakers have to be larger than 0."
1214	        # speaker embedding
1215	        if self.args.use_speaker_embedding and not self.args.use_d_vector_file:
1216	            g_src = self.emb_g(torch.from_numpy((np.array(speaker_cond_src))).unsqueeze(0)).unsqueeze(-1)

1524	            fmax=ac.mel_fmax,
1525	        )
1526	
1527	        if self.args.encoder_sample_rate:
1528	            assert batch["spec"].shape[2] == int(
1529	                batch["mel"].shape[2] / self.interpolate_factor
1530	            ), f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
1531	        else:
1532	            assert batch["spec"].shape[2] == batch["mel"].shape[2], f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
1533	

1528	            assert batch["spec"].shape[2] == int(
1529	                batch["mel"].shape[2] / self.interpolate_factor
1530	            ), f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
1531	        else:
1532	            assert batch["spec"].shape[2] == batch["mel"].shape[2], f"{batch['spec'].shape[2]}, {batch['mel'].shape[2]}"
1533	
1534	        # compute spectrogram frame lengths
1535	        batch["spec_lens"] = (batch["spec"].shape[2] * batch["waveform_rel_lens"]).int()

1535	        batch["spec_lens"] = (batch["spec"].shape[2] * batch["waveform_rel_lens"]).int()
1536	        batch["mel_lens"] = (batch["mel"].shape[2] * batch["waveform_rel_lens"]).int()
1537	
1538	        if self.args.encoder_sample_rate:
1539	            assert (batch["spec_lens"] - (batch["mel_lens"] / self.interpolate_factor).int()).sum() == 0
1540	        else:
1541	            assert (batch["spec_lens"] - batch["mel_lens"]).sum() == 0
1542	

1537	
1538	        if self.args.encoder_sample_rate:
1539	            assert (batch["spec_lens"] - (batch["mel_lens"] / self.interpolate_factor).int()).sum() == 0
1540	        else:
1541	            assert (batch["spec_lens"] - batch["mel_lens"]).sum() == 0
1542	
1543	        # zero the padding frames
1544	        batch["spec"] = batch["spec"] * sequence_mask(batch["spec_lens"]).unsqueeze(1)

1721	        self.load_state_dict(state["model"], strict=strict)
1722	
1723	        if eval:
1724	            self.eval()
1725	            assert not self.training
1726	
1727	    def load_fairseq_checkpoint(
1728	        self, config, checkpoint_dir, eval=False, strict=True

1765	        new_chk = rehash_fairseq_vits_checkpoint(checkpoint_file)
1766	        self.load_state_dict(new_chk, strict=strict)
1767	        if eval:
1768	            self.eval()
1769	            assert not self.training
1770	
1771	    @staticmethod
1772	    def init_from_config(config: "VitsConfig", samples: Union[List[List], List[Dict]] = None, verbose=True):

1781	
1782	        upsample_rate = torch.prod(torch.as_tensor(config.model_args.upsample_rates_decoder)).item()
1783	
1784	        if not config.model_args.encoder_sample_rate:
1785	            assert (
1786	                upsample_rate == config.audio.hop_length
1787	            ), f" [!] Product of upsample rates must be equal to the hop length - {upsample_rate} vs {config.audio.hop_length}"
1788	        else:
1789	            encoder_to_vocoder_upsampling_factor = config.audio.sample_rate / config.model_args.encoder_sample_rate
1790	            effective_hop_length = config.audio.hop_length * encoder_to_vocoder_upsampling_factor

1787	            ), f" [!] Product of upsample rates must be equal to the hop length - {upsample_rate} vs {config.audio.hop_length}"
1788	        else:
1789	            encoder_to_vocoder_upsampling_factor = config.audio.sample_rate / config.model_args.encoder_sample_rate
1790	            effective_hop_length = config.audio.hop_length * encoder_to_vocoder_upsampling_factor
1791	            assert (
1792	                upsample_rate == effective_hop_length
1793	            ), f" [!] Product of upsample rates must be equal to the hop length - {upsample_rate} vs {effective_hop_length}"
1794	
1795	        ap = AudioProcessor.init_from_config(config, verbose=verbose)
1796	        tokenizer, new_config = TTSTokenizer.init_from_config(config)

394	            `text_input` as text token IDs after tokenizer, `voice_samples` as samples used for cloning, `conditioning_latents`
395	            as latents used at inference.
396	
397	        """
398	        assert (
399	            "zh-cn" if language == "zh" else language in self.config.languages
400	        ), f" ❗ Language {language} is not supported. Supported languages are {self.config.languages}"
401	        # Use generally found best tuning knobs for generation.
402	        settings = {
403	            "temperature": config.temperature,

532	        for sent in text:
533	            sent = sent.strip().lower()
534	            text_tokens = torch.IntTensor(self.tokenizer.encode(sent, lang=language)).unsqueeze(0).to(self.device)
535	
536	            assert (
537	                text_tokens.shape[-1] < self.args.gpt_max_text_tokens
538	            ), " ❗ XTTS can only generate text with a maximum of 400 tokens."
539	
540	            with torch.no_grad():
541	                gpt_codes = self.gpt.generate(

640	        for sent in text:
641	            sent = sent.strip().lower()
642	            text_tokens = torch.IntTensor(self.tokenizer.encode(sent, lang=language)).unsqueeze(0).to(self.device)
643	
644	            assert (
645	                text_tokens.shape[-1] < self.args.gpt_max_text_tokens
646	            ), " ❗ XTTS can only generate text with a maximum of 400 tokens."
647	
648	            fake_inputs = self.gpt.compute_embeddings(
649	                gpt_cond_latent.to(self.device),

5	
6	
7	def _pad_data(x, length):
8	    _pad = 0
9	    assert x.ndim == 1
10	    return np.pad(x, (0, length - x.shape[0]), mode="constant", constant_values=_pad)
11	
12	

16	
17	
18	def _pad_tensor(x, length):
19	    _pad = 0.0
20	    assert x.ndim == 2
21	    x = np.pad(x, [[0, 0], [0, length - x.shape[1]]], mode="constant", constant_values=_pad)
22	    return x
23	

39	
40	    Returns:
41	        np.ndarray: Padded stop target array.
42	    """
43	    assert x.ndim == 1
44	    return np.pad(x, (0, length - x.shape[0]), mode="constant", constant_values=pad_val)
45	
46	

111	    if let_short_samples:
112	        _x_lenghts[len_diff < 0] = segment_size
113	        len_diff = _x_lenghts - segment_size
114	    else:
115	        assert all(
116	            len_diff > 0
117	        ), f" [!] At least one sample is shorter than the segment size ({segment_size}). \n {_x_lenghts}"
118	    segment_indices = (torch.rand([B]).type_as(x) * (len_diff + 1)).long()
119	    ret = segment(x, segment_indices, segment_size, pad_short=pad_short)
120	    return ret, segment_indices

86	        Returns:
87	            np.ndarray: embedding.
88	        """
89	        if self.name_to_id:
90	            return self.name_to_id[random.choices(list(self.name_to_id.keys()))[0]]
91	
92	        return None
93	

288	        embeddings = self.get_embeddings_by_name(idx)
289	        if num_samples is None:
290	            embeddings = np.stack(embeddings).mean(0)
291	        else:
292	            assert len(embeddings) >= num_samples, f" [!] {idx} has number of samples < {num_samples}"
293	            if randomize:
294	                embeddings = np.stack(random.choices(embeddings, k=num_samples)).mean(0)
295	            else:

290	            embeddings = np.stack(embeddings).mean(0)
291	        else:
292	            assert len(embeddings) >= num_samples, f" [!] {idx} has number of samples < {num_samples}"
293	            if randomize:
294	                embeddings = np.stack(random.choices(embeddings, k=num_samples)).mean(0)
295	            else:
296	                embeddings = np.stack(embeddings[:num_samples]).mean(0)
297	        return embeddings

304	        Returns:
305	            np.ndarray: embedding.
306	        """
307	        if self.embeddings:
308	            return self.embeddings[random.choices(list(self.embeddings.keys()))[0]]["embedding"]
309	
310	        return None
311	

179	                speaker_manager.load_embeddings_from_file(speakers_file)
180	            elif not c.use_d_vector_file:  # restor speaker manager with speaker ID file.
181	                speaker_ids_from_data = speaker_manager.name_to_id
182	                speaker_manager.load_ids_from_file(speakers_file)
183	                assert all(
184	                    speaker in speaker_manager.name_to_id for speaker in speaker_ids_from_data
185	                ), " [!] You cannot introduce new speakers to a pre-trained model."
186	        elif c.use_d_vector_file and c.d_vector_file:
187	            # new speaker manager with external speaker embeddings.
188	            speaker_manager.load_embeddings_from_file(c.d_vector_file)

43	
44	    x = tensors[0]
45	
46	    for t in tensors:
47	        assert torch.is_tensor(t), f"Expected torch.Tensor, got {type(t)}"
48	        assert t.device == x.device, f"Expected tensors to be on {x.device}, got {t.device}"
49	
50	        if size_range is None:

44	    x = tensors[0]
45	
46	    for t in tensors:
47	        assert torch.is_tensor(t), f"Expected torch.Tensor, got {type(t)}"
48	        assert t.device == x.device, f"Expected tensors to be on {x.device}, got {t.device}"
49	
50	        if size_range is None:
51	            assert t.size() == x.size(), f"Expected tensors with same size, got {t.size()} and {x.size()}"

47	        assert torch.is_tensor(t), f"Expected torch.Tensor, got {type(t)}"
48	        assert t.device == x.device, f"Expected tensors to be on {x.device}, got {t.device}"
49	
50	        if size_range is None:
51	            assert t.size() == x.size(), f"Expected tensors with same size, got {t.size()} and {x.size()}"
52	        else:
53	            assert (
54	                t.size()[size_range[0] : size_range[1]] == x.size()[size_range[0] : size_range[1]]

49	
50	        if size_range is None:
51	            assert t.size() == x.size(), f"Expected tensors with same size, got {t.size()} and {x.size()}"
52	        else:
53	            assert (
54	                t.size()[size_range[0] : size_range[1]] == x.size()[size_range[0] : size_range[1]]
55	            ), f"Expected tensors with same size at given dimensions, got {t.size()} and {x.size()}"
56	
57	        if dim_range[0] == dim_range[1]:
58	            assert t.dim() == dim_range[0], f"Expected number of dimensions to be {dim_range[0]}, got {t.dim()}"

54	                t.size()[size_range[0] : size_range[1]] == x.size()[size_range[0] : size_range[1]]
55	            ), f"Expected tensors with same size at given dimensions, got {t.size()} and {x.size()}"
56	
57	        if dim_range[0] == dim_range[1]:
58	            assert t.dim() == dim_range[0], f"Expected number of dimensions to be {dim_range[0]}, got {t.dim()}"
59	        elif dim_range[0] < dim_range[1]:
60	            assert (
61	                dim_range[0] <= t.dim() <= dim_range[1]

56	
57	        if dim_range[0] == dim_range[1]:
58	            assert t.dim() == dim_range[0], f"Expected number of dimensions to be {dim_range[0]}, got {t.dim()}"
59	        elif dim_range[0] < dim_range[1]:
60	            assert (
61	                dim_range[0] <= t.dim() <= dim_range[1]
62	            ), f"Expected number of dimensions to be between {dim_range[0]} and {dim_range[1]}, got {t.dim()}"
63	
64	        if data_range[0] < data_range[1]:
65	            assert data_range[0] <= t.min(), f"Expected values to be greater or equal to {data_range[0]}, got {t.min()}"

61	                dim_range[0] <= t.dim() <= dim_range[1]
62	            ), f"Expected number of dimensions to be between {dim_range[0]} and {dim_range[1]}, got {t.dim()}"
63	
64	        if data_range[0] < data_range[1]:
65	            assert data_range[0] <= t.min(), f"Expected values to be greater or equal to {data_range[0]}, got {t.min()}"
66	            assert t.max() <= data_range[1], f"Expected values to be lower or equal to {data_range[1]}, got {t.max()}"
67	
68	

62	            ), f"Expected number of dimensions to be between {dim_range[0]} and {dim_range[1]}, got {t.dim()}"
63	
64	        if data_range[0] < data_range[1]:
65	            assert data_range[0] <= t.min(), f"Expected values to be greater or equal to {data_range[0]}, got {t.min()}"
66	            assert t.max() <= data_range[1], f"Expected values to be lower or equal to {data_range[1]}, got {t.max()}"
67	
68	
69	def gaussian_filter(kernel_size: int, sigma: float) -> torch.Tensor:

124	        IEEE Transactions on Image Processing, 13, 600-612.
125	        https://ece.uwaterloo.ca/~z70wang/publications/ssim.pdf,
126	        DOI: `10.1109/TIP.2003.819861`
127	    """
128	    assert kernel_size % 2 == 1, f"Kernel size must be odd, got [{kernel_size}]"
129	    _validate_input([x, y], dim_range=(4, 5), data_range=(0, data_range))
130	
131	    x = x / float(data_range)

228	        self.kernel_size = kernel_size
229	
230	        # This check might look redundant because kernel size is checked within the ssim function anyway.
231	        # However, this check allows to fail fast when the loss is being initialised and training has not been started.
232	        assert kernel_size % 2 == 1, f"Kernel size must be odd, got [{kernel_size}]"
233	        self.kernel_sigma = kernel_sigma
234	        self.k1 = k1
235	        self.k2 = k2

185	
186	    language_name = None
187	    if language_id is not None:
188	        language = [k for k, v in model.language_manager.name_to_id.items() if v == language_id]
189	        assert len(language) == 1, "language_id must be a valid language"
190	        language_name = language[0]
191	
192	    # convert text to sequence of token IDs

290	        _vocab = [self._pad] + _vocab if self._pad is not None and len(self._pad) > 0 else _vocab
291	        self.vocab = _vocab + list(self._punctuations)
292	        if self.is_unique:
293	            duplicates = {x for x in self.vocab if self.vocab.count(x) > 1}
294	            assert (
295	                len(self.vocab) == len(self._char_to_id) == len(self._id_to_char)
296	            ), f" [!] There are duplicate characters in the character set. {duplicates}"
297	
298	    def char_to_id(self, char: str) -> int:
299	        try:

1	import logging
2	import re
3	import subprocess
4	from typing import Dict, List
5	
6	from packaging.version import Version
7	
8	from TTS.tts.utils.text.phonemizers.base import BasePhonemizer

20	espeak_version_pattern = re.compile(r"text-to-speech:\s(?P<version>\d+\.\d+(\.\d+)?)")
21	
22	
23	def get_espeak_version():
24	    output = subprocess.getoutput("espeak --version")
25	    match = espeak_version_pattern.search(output)
26	
27	    return match.group("version")

20	espeak_version_pattern = re.compile(r"text-to-speech:\s(?P<version>\d+\.\d+(\.\d+)?)")
21	
22	
23	def get_espeak_version():
24	    output = subprocess.getoutput("espeak --version")
25	    match = espeak_version_pattern.search(output)
26	
27	    return match.group("version")

27	    return match.group("version")
28	
29	
30	def get_espeakng_version():
31	    output = subprocess.getoutput("espeak-ng --version")
32	    return output.split()[3]
33	
34	

27	    return match.group("version")
28	
29	
30	def get_espeakng_version():
31	    output = subprocess.getoutput("espeak-ng --version")
32	    return output.split()[3]
33	
34	

54	    ]
55	    cmd.extend(args)
56	    logging.debug("espeakng: executing %s", repr(cmd))
57	
58	    with subprocess.Popen(
59	        cmd,
60	        stdout=subprocess.PIPE,
61	        stderr=subprocess.STDOUT,
62	    ) as p:
63	        res = iter(p.stdout.readline, b"")
64	        if not sync:
65	            p.stdout.close()

25	    Returns:
26	        np.ndarray: melspectrogram basis.
27	    """
28	    if mel_fmax is not None:
29	        assert mel_fmax <= sample_rate // 2
30	        assert mel_fmax - mel_fmin > 0
31	    return librosa.filters.mel(sr=sample_rate, n_fft=fft_size, n_mels=num_mels, fmin=mel_fmin, fmax=mel_fmax)
32	

26	        np.ndarray: melspectrogram basis.
27	    """
28	    if mel_fmax is not None:
29	        assert mel_fmax <= sample_rate // 2
30	        assert mel_fmax - mel_fmin > 0
31	    return librosa.filters.mel(sr=sample_rate, n_fft=fft_size, n_mels=num_mels, fmin=mel_fmin, fmax=mel_fmax)
32	
33	

39	    Returns:
40	        Tuple[int, int]: hop length and window length for STFT.
41	    """
42	    factor = frame_length_ms / frame_shift_ms
43	    assert (factor).is_integer(), " [!] frame_shift_ms should divide frame_length_ms"
44	    win_length = int(frame_length_ms / 1000.0 * sample_rate)
45	    hop_length = int(win_length / float(factor))
46	    return win_length, hop_length

68	
69	    Returns:
70	        np.ndarray: Decibels spectrogram.
71	    """
72	    assert (x < 0).sum() == 0, " [!] Input values must be non-negative."
73	    return gain * _log(np.maximum(1e-8, x), base)
74	
75	

128	
129	
130	def mel_to_spec(*, mel: np.ndarray = None, mel_basis: np.ndarray = None, **kwargs) -> np.ndarray:
131	    """Convert a melspectrogram to full scale spectrogram."""
132	    assert (mel < 0).sum() == 0, " [!] Input values must be non-negative."
133	    inv_mel_basis = np.linalg.pinv(mel_basis)
134	    return np.maximum(1e-10, np.dot(inv_mel_basis, mel))
135	

276	        >>> ap = AudioProcessor(**conf)
277	        >>> wav = ap.load_wav(WAV_FILE, sr=ap.sample_rate)[:5 * ap.sample_rate]
278	        >>> pitch = ap.compute_f0(wav)
279	    """
280	    assert pitch_fmax is not None, " [!] Set `pitch_fmax` before caling `compute_f0`."
281	    assert pitch_fmin is not None, " [!] Set `pitch_fmin` before caling `compute_f0`."
282	
283	    f0, voiced_mask, _ = pyin(

277	        >>> wav = ap.load_wav(WAV_FILE, sr=ap.sample_rate)[:5 * ap.sample_rate]
278	        >>> pitch = ap.compute_f0(wav)
279	    """
280	    assert pitch_fmax is not None, " [!] Set `pitch_fmax` before caling `compute_f0`."
281	    assert pitch_fmin is not None, " [!] Set `pitch_fmin` before caling `compute_f0`."
282	
283	    f0, voiced_mask, _ = pyin(
284	        y=x.astype(np.double),

399	
400	    Returns:
401	        np.ndarray: RMS normalized waveform.
402	    """
403	    assert -99 <= db_level <= 0, " [!] db_level should be between -99 and 0"
404	    wav = rms_norm(wav=x, db_level=db_level)
405	    return wav
406	

222	        else:
223	            # use stft parameters from config file
224	            self.hop_length = hop_length
225	            self.win_length = win_length
226	        assert min_level_db != 0.0, " [!] min_level_db is 0"
227	        assert (
228	            self.win_length <= self.fft_size
229	        ), f" [!] win_length cannot be larger than fft_size - {self.win_length} vs {self.fft_size}"

223	            # use stft parameters from config file
224	            self.hop_length = hop_length
225	            self.win_length = win_length
226	        assert min_level_db != 0.0, " [!] min_level_db is 0"
227	        assert (
228	            self.win_length <= self.fft_size
229	        ), f" [!] win_length cannot be larger than fft_size - {self.win_length} vs {self.fft_size}"
230	        members = vars(self)
231	        if verbose:
232	            print(" > Setting up Audio Processor...")

357	        for key in stats_config.keys():
358	            if key in skip_parameters:
359	                continue
360	            if key not in ["sample_rate", "trim_db"]:
361	                assert (
362	                    stats_config[key] == self.__dict__[key]
363	                ), f" [!] Audio param {key} does not match the value used for computing mean-var stats. {stats_config[key]} vs {self.__dict__[key]}"
364	        return mel_mean, mel_std, linear_mean, linear_std, stats_config
365	
366	    # pylint: disable=attribute-defined-outside-init

10	    return rt
11	
12	
13	def init_distributed(rank, num_gpus, group_name, dist_backend, dist_url):
14	    assert torch.cuda.is_available(), "Distributed mode requires CUDA."
15	
16	    # Set cuda device so everything is done on the right GPU.
17	    torch.cuda.set_device(rank % torch.cuda.device_count())

3	import importlib
4	import logging
5	import os
6	import re
7	import subprocess
8	import sys
9	from pathlib import Path
10	from typing import Dict

30	
31	
32	def get_git_branch():
33	    try:
34	        out = subprocess.check_output(["git", "branch"]).decode("utf8")
35	        current = next(line for line in out.split("\n") if line.startswith("*"))
36	        current.replace("* ", "")
37	    except subprocess.CalledProcessError:

30	
31	
32	def get_git_branch():
33	    try:
34	        out = subprocess.check_output(["git", "branch"]).decode("utf8")
35	        current = next(line for line in out.split("\n") if line.startswith("*"))
36	        current.replace("* ", "")
37	    except subprocess.CalledProcessError:

49	    # except:
50	    #     raise RuntimeError(
51	    #         " !! Commit before training to get the commit hash.")
52	    try:
53	        commit = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode().strip()
54	    # Not copying .git folder into docker container
55	    except (subprocess.CalledProcessError, FileNotFoundError):
56	        commit = "0000000"

49	    # except:
50	    #     raise RuntimeError(
51	    #         " !! Commit before training to get the commit hash.")
52	    try:
53	        commit = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode().strip()
54	    # Not copying .git folder into docker container
55	    except (subprocess.CalledProcessError, FileNotFoundError):
56	        commit = "0000000"

1	import os
2	import pickle as pickle_tts
3	from typing import Any, Callable, Dict, Union
4	
5	import fsspec
6	import torch
7	
8	from TTS.utils.generic_utils import get_user_data_dir

402	            # ToDo: we need a better way to handle it
403	            if "xtts" in model_name:
404	                try:
405	                    self.check_if_configs_are_equal(model_name, model_item, output_path)
406	                except:
407	                    pass
408	            else:
409	                print(f" > {model_name} is already downloaded.")
410	        else:

48	        drop_last=False,
49	        label_key="class_name",
50	    ):
51	        super().__init__(dataset_items)
52	        assert (
53	            batch_size % (num_classes_in_batch * num_gpus) == 0
54	        ), "Batch size must be divisible by number of classes times the number of data parallel devices (if enabled)."
55	
56	        label_indices = {}
57	        for idx, item in enumerate(dataset_items):

86	        self.seg = self._get_segmenter("en")
87	        self.use_cuda = use_cuda
88	        self.voice_dir = voice_dir
89	        if self.use_cuda:
90	            assert torch.cuda.is_available(), "CUDA is not availabe on this machine."
91	
92	        if tts_checkpoint:
93	            self._load_tts(tts_checkpoint, tts_config_path, use_cuda)

193	                # by cutting off from the largest duration indeces.
194	                extra_frames = dur.sum() - mel_lengths[idx]
195	                largest_idxs = torch.argsort(-dur)[:extra_frames]
196	                dur[largest_idxs] -= 1
197	                assert (
198	                    dur.sum() == mel_lengths[idx]
199	                ), f" [!] total duration {dur.sum()} vs spectrogram length {mel_lengths[idx]}"
200	                durations[idx, : text_lengths[idx]] = dur
201	
202	        # set stop targets wrt reduction factor

34	        self.out_channels = out_channels
35	        self.kernel_size = kernel_size
36	        self.n_layers = n_layers
37	        self.p_dropout = p_dropout
38	        assert n_layers > 1, "Number of layers should be larger than 0."
39	
40	        self.conv_layers = nn.ModuleList()
41	        self.norm_layers = nn.ModuleList()

103	
104	class WN(torch.nn.Module):
105	    def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
106	        super(WN, self).__init__()
107	        assert kernel_size % 2 == 1
108	        self.hidden_channels = hidden_channels
109	        self.kernel_size = (kernel_size,)
110	        self.dilation_rate = dilation_rate

346	        p_dropout=0,
347	        gin_channels=0,
348	        mean_only=False,
349	    ):
350	        assert channels % 2 == 0, "channels should be divisible by 2"
351	        super().__init__()
352	        self.channels = channels
353	        self.hidden_channels = hidden_channels

85	        :return: the waveform slices and mel spectrogram slices as lists of array slices. Index
86	        respectively the waveform and the mel spectrogram with these slices to obtain the partial
87	        utterances.
88	        """
89	        assert 0 < min_coverage <= 1
90	
91	        # Compute how many frames separate two partial utterances
92	        samples_per_frame = int((sampling_rate * mel_window_step / 1000))

91	        # Compute how many frames separate two partial utterances
92	        samples_per_frame = int((sampling_rate * mel_window_step / 1000))
93	        n_frames = int(np.ceil((n_samples + 1) / samples_per_frame))
94	        frame_step = int(np.round((sampling_rate / rate) / samples_per_frame))
95	        assert 0 < frame_step, "The rate is too high"
96	        assert frame_step <= partials_n_frames, "The rate is too low, it should be %f at least" % (
97	            sampling_rate / (samples_per_frame * partials_n_frames)
98	        )

92	        samples_per_frame = int((sampling_rate * mel_window_step / 1000))
93	        n_frames = int(np.ceil((n_samples + 1) / samples_per_frame))
94	        frame_step = int(np.round((sampling_rate / rate) / samples_per_frame))
95	        assert 0 < frame_step, "The rate is too high"
96	        assert frame_step <= partials_n_frames, "The rate is too low, it should be %f at least" % (
97	            sampling_rate / (samples_per_frame * partials_n_frames)
98	        )
99	
100	        # Compute the slices
101	        wav_slices, mel_slices = [], []

218	    if p <= 0:
219	        return module
220	
221	    # supported modules
222	    assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
223	
224	    # test whether module.weight has the right sizes wrt block_size
225	    is_conv = module.weight.ndim == 4

225	    is_conv = module.weight.ndim == 4
226	
227	    # 2D matrix
228	    if not is_conv:
229	        assert module.weight.size(1) % block_size == 0, "Input features must be a multiple of block sizes"
230	
231	    # 4D matrix
232	    else:

231	    # 4D matrix
232	    else:
233	        # 1x1 convolutions
234	        if module.kernel_size == (1, 1):
235	            assert module.in_channels % block_size == 0, "Input channels must be a multiple of block sizes"
236	        # regular convolutions
237	        else:
238	            k = module.kernel_size[0] * module.kernel_size[1]

235	            assert module.in_channels % block_size == 0, "Input channels must be a multiple of block sizes"
236	        # regular convolutions
237	        else:
238	            k = module.kernel_size[0] * module.kernel_size[1]
239	            assert k % block_size == 0, "Kernel size must be a multiple of block size"
240	
241	    def _forward_pre_hook(mod, input):
242	        # no noise for evaluation

323	
324	        self.head_dim = embed_dim // num_heads
325	        self.q_head_dim = self.head_dim
326	        self.k_head_dim = self.head_dim
327	        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
328	        self.scaling = self.head_dim**-0.5
329	
330	        self.self_attention = self_attention

329	
330	        self.self_attention = self_attention
331	        self.encoder_decoder_attention = encoder_decoder_attention
332	
333	        assert not self.self_attention or self.qkv_same_dim, (
334	            "Self-attention requires query, key and " "value to be of the same size"
335	        )
336	
337	        k_bias = True
338	        if rescale_init:

458	        is_tpu = query.device.type == "xla"
459	
460	        tgt_len, bsz, embed_dim = query.size()
461	        src_len = tgt_len
462	        assert embed_dim == self.embed_dim
463	        assert list(query.size()) == [tgt_len, bsz, embed_dim]
464	        if key is not None:
465	            src_len, key_bsz, _ = key.size()

459	
460	        tgt_len, bsz, embed_dim = query.size()
461	        src_len = tgt_len
462	        assert embed_dim == self.embed_dim
463	        assert list(query.size()) == [tgt_len, bsz, embed_dim]
464	        if key is not None:
465	            src_len, key_bsz, _ = key.size()
466	            if not torch.jit.is_scripting():

463	        assert list(query.size()) == [tgt_len, bsz, embed_dim]
464	        if key is not None:
465	            src_len, key_bsz, _ = key.size()
466	            if not torch.jit.is_scripting():
467	                assert key_bsz == bsz
468	                assert value is not None
469	                assert src_len, bsz == value.shape[:2]
470	

464	        if key is not None:
465	            src_len, key_bsz, _ = key.size()
466	            if not torch.jit.is_scripting():
467	                assert key_bsz == bsz
468	                assert value is not None
469	                assert src_len, bsz == value.shape[:2]
470	
471	        if self.has_relative_attention_bias and position_bias is None:

465	            src_len, key_bsz, _ = key.size()
466	            if not torch.jit.is_scripting():
467	                assert key_bsz == bsz
468	                assert value is not None
469	                assert src_len, bsz == value.shape[:2]
470	
471	        if self.has_relative_attention_bias and position_bias is None:
472	            position_bias = self.compute_bias(tgt_len, src_len)

480	            # treats bias in linear module as method.
481	            and not torch.jit.is_scripting()
482	            and self.q_head_dim == self.head_dim
483	        ):
484	            assert key is not None and value is not None
485	            assert attn_mask is None
486	
487	            attn_mask_rel_pos = None

481	            and not torch.jit.is_scripting()
482	            and self.q_head_dim == self.head_dim
483	        ):
484	            assert key is not None and value is not None
485	            assert attn_mask is None
486	
487	            attn_mask_rel_pos = None
488	            if position_bias is not None:

536	            if saved_state is not None and "prev_key" in saved_state:
537	                # previous time steps are cached - no need to recompute
538	                # key and value if they are static
539	                if static_kv:
540	                    assert self.encoder_decoder_attention and not self.self_attention
541	                    key = value = None
542	        else:
543	            saved_state = None

549	        elif self.encoder_decoder_attention:
550	            # encoder-decoder attention
551	            q = self.q_proj(query)
552	            if key is None:
553	                assert value is None
554	                k = v = None
555	            else:
556	                k = self.k_proj(key)

556	                k = self.k_proj(key)
557	                v = self.v_proj(key)
558	
559	        else:
560	            assert key is not None and value is not None
561	            q = self.q_proj(query)
562	            k = self.k_proj(key)
563	            v = self.v_proj(value)

563	            v = self.v_proj(value)
564	        q *= self.scaling
565	
566	        if self.bias_k is not None:
567	            assert self.bias_v is not None
568	            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
569	            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
570	            if attn_mask is not None:

587	        if saved_state is not None:
588	            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
589	            if "prev_key" in saved_state:
590	                _prev_key = saved_state["prev_key"]
591	                assert _prev_key is not None
592	                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
593	                if static_kv:
594	                    k = prev_key

592	                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
593	                if static_kv:
594	                    k = prev_key
595	                else:
596	                    assert k is not None
597	                    k = torch.cat([prev_key, k], dim=1)
598	                src_len = k.size(1)
599	            if "prev_value" in saved_state:

597	                    k = torch.cat([prev_key, k], dim=1)
598	                src_len = k.size(1)
599	            if "prev_value" in saved_state:
600	                _prev_value = saved_state["prev_value"]
601	                assert _prev_value is not None
602	                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
603	                if static_kv:
604	                    v = prev_value

602	                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
603	                if static_kv:
604	                    v = prev_value
605	                else:
606	                    assert v is not None
607	                    v = torch.cat([prev_value, v], dim=1)
608	            prev_key_padding_mask: Optional[Tensor] = None
609	            if "prev_key_padding_mask" in saved_state:

607	                    v = torch.cat([prev_value, v], dim=1)
608	            prev_key_padding_mask: Optional[Tensor] = None
609	            if "prev_key_padding_mask" in saved_state:
610	                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
611	            assert k is not None and v is not None
612	            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
613	                key_padding_mask=key_padding_mask,
614	                prev_key_padding_mask=prev_key_padding_mask,

620	            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
621	            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
622	            saved_state["prev_key_padding_mask"] = key_padding_mask
623	            # In this branch incremental_state is never None
624	            assert incremental_state is not None
625	            incremental_state = self._set_input_buffer(incremental_state, saved_state)
626	        assert k is not None
627	        assert k.size(1) == src_len

622	            saved_state["prev_key_padding_mask"] = key_padding_mask
623	            # In this branch incremental_state is never None
624	            assert incremental_state is not None
625	            incremental_state = self._set_input_buffer(incremental_state, saved_state)
626	        assert k is not None
627	        assert k.size(1) == src_len
628	
629	        # This is part of a workaround to get around fork/join parallelism

623	            # In this branch incremental_state is never None
624	            assert incremental_state is not None
625	            incremental_state = self._set_input_buffer(incremental_state, saved_state)
626	        assert k is not None
627	        assert k.size(1) == src_len
628	
629	        # This is part of a workaround to get around fork/join parallelism
630	        # not supporting Optional types.

631	        if key_padding_mask is not None and key_padding_mask.dim() == 0:
632	            key_padding_mask = None
633	
634	        if key_padding_mask is not None:
635	            assert key_padding_mask.size(0) == bsz
636	            assert key_padding_mask.size(1) == src_len
637	
638	        if self.add_zero_attn:

632	            key_padding_mask = None
633	
634	        if key_padding_mask is not None:
635	            assert key_padding_mask.size(0) == bsz
636	            assert key_padding_mask.size(1) == src_len
637	
638	        if self.add_zero_attn:
639	            assert v is not None

635	            assert key_padding_mask.size(0) == bsz
636	            assert key_padding_mask.size(1) == src_len
637	
638	        if self.add_zero_attn:
639	            assert v is not None
640	            src_len += 1
641	            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
642	            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)

653	
654	        attn_weights = torch.bmm(q, k.transpose(1, 2))
655	        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
656	
657	        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
658	
659	        if attn_mask is not None:
660	            attn_mask = attn_mask.unsqueeze(0)

694	        attn_weights_float = F.softmax(attn_weights, dim=-1)
695	        attn_weights = attn_weights_float.type_as(attn_weights)
696	        attn_probs = self.dropout_module(attn_weights)
697	
698	        assert v is not None
699	        attn = torch.bmm(attn_probs, v)
700	        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
701	        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)

696	        attn_probs = self.dropout_module(attn_weights)
697	
698	        assert v is not None
699	        attn = torch.bmm(attn_probs, v)
700	        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
701	        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
702	        attn = self.out_proj(attn)
703	        attn_weights: Optional[Tensor] = None

370	        conv_type: str = "default",
371	    ):
372	        super().__init__()
373	
374	        assert mode in {"default", "layer_norm"}
375	
376	        def block(
377	            n_in,