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}")

"""

Class to sample from a trained model.

"""

def __init__(

self,

config: InferenceConfig,

tokenizer_cls: Type[TrainedBPETokeniser],

decoder_cls: Type[Decoder],

data_adapter_fn,

use_kv_cache: Optional[Literal["none", "flash_decoding", "vanilla"]] = 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

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=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":

for block in self.model.transformer.h:

if block.attn.attn_kernel_type != "fa2":

raise Exception(

f"kv_cache only supported for flash attention 2 but found {block.attn.attn_kernel_type} inside model!"

)

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,

)

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,

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"),

"""

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")

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)

return cache_path

else:

if os.path.exists(file_or_uri):

return file_or_uri

else:

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)

texts: list[str],

spk_cond_paths: list[Optional[str]],

spkemb_model,

first_stage_model,

second_stage_model,

enhancer: Optional[Literal["df"] | BaseEnhancer],

first_stage_ckpt_path: str,

second_stage_ckpt_path: str,

guidance_scale: Optional[float],

max_new_tokens: int,

top_k: Optional[int],

top_p: Optional[float],

temperature: Optional[float],

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")

save_result_metadata(

wav_file,

ref_name,

text,

first_stage_ckpt_path,

second_stage_ckpt_path,

)

text: str,

spk_cond_path: Optional[str],

spkemb_model,

first_stage_model,

second_stage_model,

enhancer: Optional[Literal["df"] | BaseEnhancer],

first_stage_ckpt_path: str,

second_stage_ckpt_path: str,

guidance_scale: Optional[float],

max_new_tokens: int,

top_k: Optional[int],

top_p: Optional[float],

temperature: Optional[float],

) -> 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,

smodel = SpeakerEncoder(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

)

"""

Sample from a trained model.

"""

huggingface_repo_id: str

"""Absolute path to the model directory."""

spk_cond_path: str

"""Path to speaker reference file. Supports: wav, flac & mp3"""

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

"""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"] = "bfloat16"

"""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[Literal["flash_decoding", "vanilla"]] = 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 flash attention 2 with hand implemented kv-cache."""

output_dir: str = "samples/"

"""Relative path to output directory"""

guidance_scale: Optional[float] = None

"""Guidance scale for sampling."""

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