from fam.llm.fast_inference import TTS import string import json from glob import glob import torch import os import torchaudio import subprocess import shutil import soundfile as sf import pyloudnorm as pyln import noisereduce as nr from moviepy import * from pydub import AudioSegment from fam.llm.adapters import FlattenedInterleavedEncodec2Codebook, TiltedEncodec from fam.llm.decoders import Decoder, EncodecDecoder from fam.llm.enhancers import BaseEnhancer, get_enhancer from fam.llm.model import GPT, GPTConfig from fam.llm.utils import ( check_audio_file, get_default_dtype, # get_default_use_kv_cache, normalize_text, ) from fam.quantiser.audio.speaker_encoder.model import SpeakerEncoder from fam.quantiser.text.tokenise import TrainedBPETokeniser import tyro from huggingface_hub import snapshot_download from typing import List, Literal, Optional, Tuple, Type, Union import dataclasses import hashlib import json import os import pathlib from contextlib import nullcontext from dataclasses import dataclass import tqdm import tqdm.contrib.concurrent import tempfile import textwrap def remove_punctuation(sentence): translator = str.maketrans('', '', string.punctuation) sentence = sentence.translate(translator) # Remove line breaks sentence = sentence.replace('\n', ' ').replace('\r', '') return sentence # def run_audio_generation_v2(new_text,accent='None'): # tts = TTS() # new_text = new_text.replace('\n', ' ').replace('\r', '') # new_text_mod = remove_punctuation(new_text) # new_text_split = new_text_mod.split() # for word in new_text_split: # if len(word)>=2 and word.isupper(): # new_text = new_text.replace(word, " ".join([*word])) # wav_file = tts.synthesise( # text=new_text, # spk_ref_path="./tmp/audio/speaker_wav.wav" # you can use any speaker reference file (WAV, OGG, MP3, FLAC, etc.) # ) # sf.write('audio/output.wav', wav_file, samplerate=22050) def run_audio_generation_v2(new_text, accent=None): # check for abbreviations in new text. need to add - after each letter so that audio comes out okay new_text = new_text.replace('\n', ' ').replace('\r', '') new_text_mod = remove_punctuation(new_text) new_text_split = new_text_mod.split() for word in new_text_split: if len(word)>=2 and word.isupper(): new_text = new_text.replace(word, " ".join([*word])) print(new_text) if len(new_text)<=220: sampling_config = SamplingControllerConfig(spk_cond_path="./tmp/audio/input_src/0.wav", text=new_text, output_dir='./tmp/audio/', output_name='generated-custom.wav') metavoice_gen(sampling_config) else: new_texts = new_text.split('. ') #textwrap.wrap(new_text, 220) new_texts = [txt +"." for txt in new_texts] output_names = [] for idx, new_text in enumerate(new_texts): output_name = "-{}.".format(idx).join('generated-custom.wav'.split('.')) output_names.append(output_name) sampling_config = SamplingControllerConfig(spk_cond_path="./tmp/audio/input_src/0.wav", text=new_text, output_dir='./tmp/audio/multiple/', output_name=output_name) metavoice_gen(sampling_config) #audio_files = ['./tmp/audio/multiple/'+'/'+ aud for aud in os.listdir('./tmp/audio/multiple/') if aud.endswith(".wav")] audio_files = ['./tmp/audio/multiple/'+'/'+ aud for aud in output_names] print(audio_files) clips = [(AudioFileClip(clip)) for clip in audio_files] final_clip = concatenate_audioclips(clips) final_clip.write_audiofile('./tmp/audio/generated-custom.wav') # adjust loudness data, rate = sf.read("./tmp/audio/input_audio.wav") # load audio (with shape (samples, channels)) meter = pyln.Meter(rate) # create BS.1770 meter loudness_target = meter.integrated_loudness(data) # measure loudness mod_data, mod_rate = sf.read("./tmp/audio/generated-custom.wav") # load audio (with shape (samples, channels)) mod_meter = pyln.Meter(mod_rate) # create BS.1770 meter loudness_gen = mod_meter.integrated_loudness(mod_data) # measure loudness loudness_normalized_gen = pyln.normalize.loudness(mod_data, loudness_gen, loudness_target) sf.write('./tmp/audio/generated-custom.wav', loudness_normalized_gen, mod_rate) @dataclass class InferenceConfig: ckpt_path: str # path to checkpoint output_dir: str num_samples: int = 10 # number of samples to draw seed: int = 1337 # random seed device: str = "cuda" dtype: str = "bfloat16" compile: bool = False init_from: str = "resume" # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl') def __str__(self): field_strs = [] for field in dataclasses.fields(self): value = getattr(self, field.name) field_strs.append(f" {field.name}: {value}") return "InferenceConfig:\n" + "\n".join(field_strs) class Model: def __init__( self, config: InferenceConfig, tokenizer_cls: Type[TrainedBPETokeniser], decoder_cls: Type[Decoder], data_adapter_fn, use_kv_cache= None, ): # TODO: disentangle the encodec stuff and numbers etc with rest of this code (esp at encoder-only / second stage model inference) # TODO: remove magic number self._encodec_codes_pad_token = 1024 self._num_encodec_codebooks = 8 self.config = config self.use_kv_cache = use_kv_cache torch.manual_seed(config.seed) torch.cuda.manual_seed(config.seed) torch.backends.cuda.matmul.allow_tf32 = True if config.dtype != "float32" else False # allow tf32 on matmul torch.backends.cudnn.allow_tf32 = True if config.dtype != "float32" else False # allow tf32 on cudnn device_type = "cuda" if "cuda" in config.device else "cpu" # for later use in torch.autocast self.ptdtype = { "float32": torch.float32, "tfloat32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16, }[config.dtype] self._ctx = ( nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type=device_type, dtype=self.ptdtype) ) self.use_bpe_tokenizer = False self.load_meta = None self.speaker_cond = None self.meta = None self.model = None self.checkpoint_config = None self.vocab_sizes = None self.smodel = None self._init_model() self.tokenizer = tokenizer_cls(**self.meta["tokenizer"]) self.decoder = decoder_cls( tokeniser_decode_fn=self.tokenizer.decode, output_dir=self.config.output_dir, data_adapter_fn=data_adapter_fn, ) def _init_model(self): if self.config.init_from == "resume": # init from a model saved in a specific directory checkpoint = torch.load(self.config.ckpt_path, map_location=self.config.device) self.vocab_sizes = checkpoint["model_args"]["vocab_sizes"] self.load_meta = False self.speaker_cond = False if "config" in checkpoint: self.checkpoint_config = checkpoint["config"] self.meta = checkpoint["meta"] load_meta = True if load_meta: self.use_bpe_tokenizer = "stoi" not in self.meta or "itos" not in self.meta self.speaker_cond = self.meta.get("speaker_cond") if self.speaker_cond: speaker_emb_size = self.meta["speaker_emb_size"] model_args = checkpoint["model_args"] if "causal" in self.checkpoint_config and self.checkpoint_config["causal"] is False: self._encodec_ctx_window = model_args["block_size"] gptconf = GPTConfig(**model_args) # TODO: rename `speaker_emb_dim` to `speaker_emb_size`. self.model = GPT(gptconf, speaker_emb_dim=speaker_emb_size if self.speaker_cond else None) state_dict = checkpoint["model"] unwanted_prefix = "_orig_mod." for k, v in list(state_dict.items()): if k.startswith(unwanted_prefix): state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k) self.model.load_state_dict(state_dict) # model self.model.eval() self.model.to(self.config.device) if self.config.compile: from einops._torch_specific import allow_ops_in_compiled_graph allow_ops_in_compiled_graph() self.model = torch.compile(self.model) # type: ignore if self.use_kv_cache is not None: if "causal" in self.checkpoint_config and self.checkpoint_config["causal"] is False: raise Exception("kv_cache not supported for non-causal models!") if self.use_kv_cache == "flash_decoding": self.model.enable_kv_cache() for block in self.model.transformer.h: block.attn.attn_kernel_type = "fd" elif self.use_kv_cache == "vanilla": self.model.enable_kv_cache() else: raise NotImplementedError(f"kv_cache type {self.use_kv_cache} not implemented!") def causal_sample( self, *, texts: list[str], batch_size: int, max_new_tokens: int, temperature: Optional[float], top_k: Optional[int], top_p: Optional[float], speaker_embs: Optional[torch.Tensor] = None, guidance_scale: Optional[float] = None, ) -> list[torch.Tensor]: """ Returns list of torch.Tensors of tokens. Each tensor is of shape (1, c, t) where c is the number of codebooks. Any flattening / inteleaving / tilting gets reversed before the output is returned. """ if speaker_embs is not None: assert len(texts) == len(speaker_embs) encoded_texts = [self.tokenizer.encode(text) for text in texts] ## create multiple hierarchies and get seq_lens seq_lens = [] xs = [] for i, encoded_text in enumerate(encoded_texts): encoded_text = torch.tensor([encoded_text], dtype=torch.long, device=self.config.device) # TODO: remove magic number xs.append( torch.cat( # [1st hierarchy of text, *remaining hierarchies of padded tokens] # TODO: self.vocab_sizes should be from the model config? [encoded_text, *[torch.ones_like(encoded_text) * 1024] * (len(self.vocab_sizes) - 1)], dim=0, ).unsqueeze(0) ) # b x [(b=1, c, t)] seq_lens.append(xs[-1].shape[-1]) max_len = max(seq_lens) assert len(xs) == len(seq_lens) ## equalise the shapes in the batch. we can use torch.zeros as tokens > seq_lens will be masked out. x = torch.zeros((len(encoded_texts), xs[0].shape[1], max_len), dtype=torch.long, device=self.config.device) for i, _xs in enumerate(xs): assert _xs.shape[-1] == seq_lens[i] x[i, :, : seq_lens[i]] = _xs ## check that the input is correct for i in range(x.shape[0]): assert x[i, 0, : seq_lens[i]].tolist() == encoded_texts[i] # TODO: remove magic number if x.shape[1] > 1: assert set(x[i, 1, : seq_lens[i]].tolist()) == set([1024]) assert x.shape[0] == speaker_embs.shape[0] if speaker_embs is not None else True if self.speaker_cond is False: speaker_embs = None # run sampling loop with torch.no_grad(): with self._ctx: # type: ignore to_return = [] for k in range(self.config.num_samples): assert seq_lens is not None assert batch_size is not None if max(seq_lens) + max_new_tokens >= self.model.config.block_size: raise Exception( f"max_new_tokens {max_new_tokens} too large! Choose {self.model.config.block_size - max(seq_lens) - 1} instead." ) y = self.model.generate( x, max_new_tokens, seq_lens=seq_lens, temperature=temperature, top_k=top_k, top_p=top_p, speaker_embs=speaker_embs, batch_size=batch_size, guidance_scale=guidance_scale, dtype=self.ptdtype, end_of_audio_token=self.tokenizer.offset - 1, end_of_text_token=self.tokenizer.eot_token, ) for i in range(len(y)): to_return.append(self.decoder.decode(tokens=y[i].tolist(), causal=True)) return to_return def non_causal_sample( self, *, texts: list[str], encodec_tokens: list[torch.Tensor], batch_size: int, top_k: Optional[int], temperature: Optional[float], speaker_embs: Optional[torch.Tensor] = None, ) -> list[str]: """ Returns paths to saved audio files. """ if speaker_embs is not None: assert len(texts) == len(speaker_embs) encoded_texts = [self.tokenizer.encode(text) for text in texts] # setup input # TODO: same code is used during data prep. refactor padded_hierarchies_inputs = [] for encoded_text, encodec_token in zip(encoded_texts, encodec_tokens): x = torch.tensor(encoded_text, dtype=torch.long, device=self.config.device)[ None, None, ... ] # (b=1, c=1, t) # TODO: should only happen if decoder is encodecdeocder? assert encodec_token.shape[0] == 1 encodec_token = encodec_token[0].tolist() # (b=1, c, t) -> (c, t) assert len(encodec_token) >= 1 and len(encodec_token) <= self._num_encodec_codebooks ## setup hierarchies of tokens # TODO: refactor and merge with code in processing.py text_tokens = encoded_text # (t,) hierarchies_in = [] hierarchies_in.append(text_tokens + encodec_token[0] + [self._encodec_codes_pad_token]) hierarchies_in.append( [self._encodec_codes_pad_token] * len(text_tokens) + encodec_token[1] + [self._encodec_codes_pad_token] ) ## adding padding / cutting to the right size as needed # TODO: refactor and merge with code in processing.py padded_hierarchies_input = [] for _, t_hierarchy in enumerate(hierarchies_in): assert len(t_hierarchy) == len(hierarchies_in[0]) if len(t_hierarchy) < self._encodec_ctx_window: padded_hierarchies_input.append( t_hierarchy + [self._encodec_codes_pad_token] * (self._encodec_ctx_window - len(t_hierarchy)) ) elif len(t_hierarchy) > self._encodec_ctx_window: padded_hierarchies_input.append(t_hierarchy[: self._encodec_ctx_window]) else: padded_hierarchies_input.append(t_hierarchy) padded_hierarchies_inputs.append(padded_hierarchies_input) ## check that the input is correct in_x = torch.tensor(padded_hierarchies_inputs, dtype=torch.long, device=self.config.device) assert in_x.shape[0] == speaker_embs.shape[0] if speaker_embs is not None else True if self.speaker_cond is False: speaker_embs = None # run sampling loop with torch.no_grad(): with self._ctx: # type: ignore to_return = [] for k in range(self.config.num_samples): y = self.model.generate( in_x, None, temperature=temperature, top_k=top_k, # TODO: handle separate top_p for this model explicitly top_p=None, speaker_embs=speaker_embs, batch_size=batch_size, guidance_scale=None, ) b_tokens = torch.cat([in_x, y], dim=1) for tokens in b_tokens: try: to_return.append(self.decoder.decode(tokens=tokens.tolist(), causal=False)) except Exception as e: print("failed to run MBD.") print(f"reason: {str(e)}") to_return.append(None) return to_return def __call__( self, *, texts: list[str], batch_size: int, max_new_tokens: Optional[int], top_k: Optional[int], top_p: Optional[float], temperature: Optional[float], encodec_tokens: Optional[list[torch.Tensor]] = None, speaker_embs: Optional[torch.Tensor] = None, guidance_scale: Optional[float] = None, ): if self.checkpoint_config.get("causal", True): return self.causal_sample( texts=texts, batch_size=batch_size, speaker_embs=speaker_embs, guidance_scale=guidance_scale, max_new_tokens=max_new_tokens, top_k=top_k, top_p=top_p, temperature=temperature, ) else: assert encodec_tokens is not None assert guidance_scale is None assert max_new_tokens is None assert top_p is None return self.non_causal_sample( texts=texts, encodec_tokens=encodec_tokens, batch_size=batch_size, speaker_embs=speaker_embs, top_k=top_k, temperature=temperature, ) def save_result_metadata(wav_path, ref_path, text, first_stage_ckpt_path, second_stage_ckpt_path): if first_stage_ckpt_path is None or second_stage_ckpt_path is None: return json.dump( { "speaker": ref_path, "text": text, }, pathlib.Path(str(wav_path) + ".json").open("w"), ) def get_cached_file(file_or_uri: str): """ If it's an s3 file, download it to a local temporary file and return that path. Otherwise return the path as is. """ is_uri = file_or_uri.startswith("http") cache_path = None if is_uri: ext = pathlib.Path(file_or_uri).suffix # hash the file path to get the cache name _cache_name = "audio_" + hashlib.md5(file_or_uri.encode("utf-8")).hexdigest() + ext os.makedirs(os.path.expanduser("~/.cache/fam/"), exist_ok=True) cache_path = os.path.expanduser(f"~/.cache/fam/{_cache_name}") if not os.path.exists(cache_path): command = f"curl -o {cache_path} {file_or_uri}" subprocess.run(command, shell=True, check=True) else: if os.path.exists(file_or_uri): cache_path = file_or_uri else: raise FileNotFoundError(f"File {file_or_uri} not found!") return cache_path def get_cached_embedding(local_file_path: str, spkemb_model): if not os.path.exists(local_file_path): raise FileNotFoundError(f"File {local_file_path} not found!") # hash the file path to get the cache name _cache_name = "embedding_" + hashlib.md5(local_file_path.encode("utf-8")).hexdigest() + ".pt" os.makedirs(os.path.expanduser("~/.cache/fam/"), exist_ok=True) cache_path = os.path.expanduser(f"~/.cache/fam/{_cache_name}") if not os.path.exists(cache_path): spk_emb = spkemb_model.embed_utterance_from_file(local_file_path, numpy=False).unsqueeze(0) # (b=1, c) torch.save(spk_emb, cache_path) else: spk_emb = torch.load(cache_path) return spk_emb def _sample_utterance_batch( texts: list[str], spk_cond_paths: list[Optional[str]], spkemb_model, first_stage_model, second_stage_model, enhancer: Optional[Union[Literal["df"], BaseEnhancer]], first_stage_ckpt_path: str, second_stage_ckpt_path: str, guidance_scale: Optional[Tuple[float, float]], max_new_tokens: int, top_k: Optional[int], top_p: Optional[float], temperature: Optional[float], output_name: str, output_dir: str, batch_size: int = 128, ) -> List[str]: speaker_embs = [] refs = spk_cond_paths.copy() # multithreaded loop to cache all the files spk_cond_paths = tqdm.contrib.concurrent.thread_map( get_cached_file, spk_cond_paths, desc="getting cached speaker ref files" ) for i, (text, spk_cond_path) in tqdm.tqdm( enumerate(zip(texts, spk_cond_paths)), total=len(texts), desc="calculating speaker embeddings" ): texts[i] = normalize_text(text) speaker_embs.append(get_cached_embedding(spk_cond_path, spkemb_model) if spk_cond_path else None) b_speaker_embs = torch.cat(speaker_embs, dim=0) b_tokens = first_stage_model( texts=texts, speaker_embs=b_speaker_embs, batch_size=batch_size, guidance_scale=guidance_scale, top_p=top_p, top_k=top_k, temperature=temperature, max_new_tokens=max_new_tokens, ) # TODO: set batch size for second stage model! wav_files = second_stage_model( texts=texts, encodec_tokens=b_tokens, speaker_embs=b_speaker_embs, batch_size=batch_size, guidance_scale=None, top_p=None, top_k=top_k, temperature=temperature, max_new_tokens=None, ) for text, tokens, speaker_embs, ref_name, wav_file in zip(texts, b_tokens, b_speaker_embs, refs, wav_files): if wav_file is None: continue with tempfile.NamedTemporaryFile(suffix=".wav") as enhanced_tmp: if enhancer is not None: enhancer = get_enhancer(enhancer) if isinstance(enhancer, str) else enhancer enhancer(str(wav_file) + ".wav", enhanced_tmp.name) # copy enhanced_tmp.name back to wav_file print(f"copying enhanced file from {enhanced_tmp.name} to {str(wav_file) + '.wav'}.") shutil.copy2(enhanced_tmp.name, str(wav_file) + ".wav") shutil.copy2(str(wav_file) + ".wav", os.path.join(output_dir, output_name)) save_result_metadata( wav_file, ref_name, text, first_stage_ckpt_path, second_stage_ckpt_path, ) return [str(w) + ".wav" if not str(w).endswith(".wav") else str(w) for w in wav_files] def sample_utterance( text: str, spk_cond_path: Optional[str], spkemb_model, first_stage_model, second_stage_model, enhancer: Optional[Union[Literal["df"], BaseEnhancer]], first_stage_ckpt_path: str, second_stage_ckpt_path: str, guidance_scale: Optional[Tuple[float, float]], max_new_tokens: int, top_k: Optional[int], top_p: Optional[float], temperature: Optional[float], output_name: str, output_dir: str, ) -> str: # NOTE: supports max. 220 characters atm. # Long form synthesis coming soon... MAX_CHARS = 220 if len(text) > MAX_CHARS: print( f"\n***WARNING: Max {MAX_CHARS} characters supported. Provided: {len(text)}. Truncating and generating speech...Can lead to unpredictable speech at the end.***" ) return _sample_utterance_batch( texts=[text], spk_cond_paths=[spk_cond_path], spkemb_model=spkemb_model, first_stage_model=first_stage_model, second_stage_model=second_stage_model, enhancer=enhancer, first_stage_ckpt_path=first_stage_ckpt_path, second_stage_ckpt_path=second_stage_ckpt_path, batch_size=1, guidance_scale=guidance_scale, max_new_tokens=max_new_tokens, top_k=top_k, top_p=top_p, temperature=temperature, output_name = output_name, output_dir = output_dir )[0] def build_models(config_first_stage, config_second_stage, model_dir, device, use_kv_cache): smodel = SpeakerEncoder( weights_fpath=os.path.join(model_dir, "speaker_encoder.pt"), device=device, eval=True, verbose=False ) data_adapter = FlattenedInterleavedEncodec2Codebook(end_of_audio_token=1024) llm_first_stage = Model( config_first_stage, TrainedBPETokeniser, EncodecDecoder, data_adapter_fn=data_adapter.decode, use_kv_cache=use_kv_cache, ) data_adapter_second_stage = TiltedEncodec(end_of_audio_token=1024) llm_second_stage = Model( config_second_stage, TrainedBPETokeniser, EncodecDecoder, data_adapter_fn=data_adapter_second_stage.decode ) return smodel, llm_first_stage, llm_second_stage def get_first_stage_path(model_dir: str): """Absolute path to checkpoint for the first stage model.""" return os.path.join(os.path.expanduser(model_dir), "first_stage.pt") def get_second_stage_path(model_dir: str): """Absolute path to checkpoint for the second stage model.""" return os.path.join(os.path.expanduser(model_dir), "second_stage.pt") @dataclass class SamplingControllerConfig: """ Sample from a trained model. """ spk_cond_path: str """Path to speaker reference file. Min. 30s of audio required. Supports both local paths & public URIs. Audio formats: wav, flac & mp3""" huggingface_repo_id: str = "metavoiceio/metavoice-1B-v0.1" """Absolute path to the model directory.""" text: str = ( "This is a demo of text to speech by MetaVoice-1B, an open-source foundational audio model by MetaVoice." ) """Text to synthesise.""" num_samples: int = 1 """Number of samples to generate from each model.""" max_new_tokens: int = 864 * 2 """Maximum number of new tokens to generate from the first stage model.""" temperature: float = 1.0 """Temperature for sampling applied to both models.""" top_k: Optional[int] = 200 """Top k for sampling applied to both models.""" top_p: Optional[float] = None """Top p for sampling applied to first-stage model.""" seed: int = 1337 """Random seed for sampling.""" device: Literal["cuda", "cpu"] = "cuda" """Device to use for sampling.""" dtype: Literal["bfloat16", "float16", "float32", "tfloat32"] = get_default_dtype() """Data type to use for sampling.""" compile: bool = False """Whether to compile the model using PyTorch 2.0.""" enhancer: Optional[Literal["df"]] = "df" """Enhancer to use for post-processing.""" init_from: str = "resume" """Either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl').""" use_kv_cache: Optional[None] = None # """Type of kv caching to use for inference: 1) [none] no kv caching, 2) [flash_decoding] use the # flash decoding kernel, 3) [vanilla] use torch attention with hand implemented kv-cache.""" output_dir: str = "samples/" """Relative path to output directory""" guidance_scale: Optional[Tuple[float, float]] = (3.0, 1.0) """Guidance scale for sampling: (speaker conditioning guidance_scale, prompt conditioning guidance scale).""" batch_size: int = 128 """Batch size to use for sampling. Note that the batch size gets doubled when guidance is used. For H100, and 1B model, 1 w/ guidance and 1 w/o guidance work well (without kv-caching). With kv-caching, 128 (w/o guidance) and 64 (w/ guidance) works well.""" output_name:str = "generated-custom.wav" def metavoice_gen(sampling_config): sampling_config = sampling_config #tyro.cli(SamplingControllerConfig, use_underscores=True) check_audio_file(sampling_config.spk_cond_path) model_dir = snapshot_download(repo_id=sampling_config.huggingface_repo_id) first_stage_ckpt_path = get_first_stage_path(model_dir) second_stage_ckpt_path = get_second_stage_path(model_dir) config_first_stage = InferenceConfig( ckpt_path=first_stage_ckpt_path, num_samples=sampling_config.num_samples, seed=sampling_config.seed, device=sampling_config.device, dtype=sampling_config.dtype, compile=sampling_config.compile, init_from=sampling_config.init_from, output_dir=sampling_config.output_dir, ) config_second_stage = InferenceConfig( ckpt_path=second_stage_ckpt_path, num_samples=sampling_config.num_samples, seed=sampling_config.seed, device=sampling_config.device, dtype=sampling_config.dtype, compile=sampling_config.compile, init_from=sampling_config.init_from, output_dir=sampling_config.output_dir, ) # sampling_config.max_new_tokens *= ( # 2 # deal with max_new_tokens for flattened interleaving! (should scale with num_codebooks?) # ) # define models smodel, llm_first_stage, llm_second_stage = build_models( config_first_stage, config_second_stage, model_dir=model_dir, device=sampling_config.device, use_kv_cache=sampling_config.use_kv_cache, ) sample_utterance( sampling_config.text, os.path.expanduser(sampling_config.spk_cond_path), smodel, llm_first_stage, llm_second_stage, sampling_config.enhancer, first_stage_ckpt_path, second_stage_ckpt_path, sampling_config.guidance_scale, max_new_tokens=sampling_config.max_new_tokens, top_k=sampling_config.top_k, top_p=sampling_config.top_p, temperature=sampling_config.temperature, output_name = sampling_config.output_name, output_dir=sampling_config.output_dir, )