Add model inference code

README.md

@@ -9,4 +9,14 @@ library_name: metavoice
 inference: false
 pipeline_tag: text-to-speech
 ---
-Converted safetensors version of [metavoiceio/metavoice-1B-v0.1](https://huggingface.co/metavoiceio/metavoice-1B-v0.1)
+Converted safetensors version of [metavoiceio/metavoice-1B-v0.1](https://huggingface.co/metavoiceio/metavoice-1B-v0.1)
+
+# Usage
+1. Install the original model code by running metavoice.sh
+2. Run the following code:
+
+```python
+from metavoice import MetaVoiceModel
+model = MetaVoiceModel("sheepymeh/metavoice-1B-v0.1")
+model.generate("Hello world!")
+```

metavoice.py

@@ -0,0 +1,234 @@
+import fam.llm.fast_inference_utils
+from fam.llm.fast_inference import TTS as FAMTTS
+from fam.llm.inference import Model as FAMModel
+from fam.llm.inference import InferenceConfig
+from fam.llm.adapters.tilted_encodec import TiltedEncodec
+from fam.llm.adapters.flattened_encodec import FlattenedInterleavedEncodec2Codebook
+from fam.llm.decoders import EncodecDecoder
+from fam.llm.enhancers import get_enhancer
+from fam.llm.utils import get_default_dtype, get_device
+from fam.llm.fast_model import Transformer
+from fam.llm.model import GPT, GPTConfig
+from fam.quantiser.text.tokenise import TrainedBPETokeniser
+from fam.quantiser.audio.speaker_encoder.model import SpeakerEncoder as FAMSpeakerEncoder
+from fam.quantiser.audio.speaker_encoder.model import mel_n_channels, model_hidden_size, model_embedding_size, model_num_layers
+
+import os
+from pathlib import Path
+from typing import Optional, Union
+from json import load, dump
+from base64 import b64encode, b64decode
+
+import torch
+from torch import nn
+from huggingface_hub import snapshot_download, HfFileSystem
+from safetensors.torch import load_model, save_model
+
+def convert_to_safetensors(
+	stage1_path: str,
+	stage2_path: str,
+	spk_emb_ckpt_path: str,
+	precision: torch.dtype,
+	output_path: str
+):
+	config_second_stage = InferenceConfig(
+		ckpt_path=stage2_path,
+		num_samples=1,
+		seed=0,
+		device='cpu',
+		dtype='float16' if precision == torch.float16 else 'bfloat16',
+		compile=False,
+		init_from='resume',
+		output_dir='.',
+	)
+	data_adapter_second_stage = TiltedEncodec(end_of_audio_token=512)
+	stage2_model = Model(config_second_stage, TrainedBPETokeniser, EncodecDecoder, data_adapter_fn=data_adapter_second_stage.decode)
+
+	stage2_checkpoint = torch.load(stage2_path, map_location='cpu')
+	stage2_state_dict = stage2_checkpoint['model']
+	unwanted_prefix = '_orig_mod.'
+	for k in stage2_state_dict.keys():
+		if k.startswith(unwanted_prefix):
+			stage2_state_dict[k[len(unwanted_prefix) :]] = stage2_state_dict.pop(k)
+	save_model(stage2_model.model, os.path.join(output_path, 'second_stage.safetensors'))
+
+	stage1_model, tokenizer, smodel = fam.llm.fast_inference_utils._load_model(stage1_path, spk_emb_ckpt_path, 'cpu', precision)
+	tokenizer_info = torch.load(stage1_path, map_location='cpu').get('meta', {}).get('tokenizer', {})
+	save_model(stage1_model, os.path.join(output_path, 'first_stage.safetensors'))
+	save_model(smodel, os.path.join(output_path, 'speaker_encoder.safetensors'))
+
+	with open(os.path.join(output_path, 'config.json'), 'w') as f:
+		tokenizer_info['mergeable_ranks'] = {b64encode(k).decode('ascii'): v for k, v in tokenizer_info['mergeable_ranks'].items()}
+		stage2_checkpoint['meta']['tokenizer']['mergeable_ranks'] = {b64encode(k).decode('ascii'): v for k, v in stage2_checkpoint['meta']['tokenizer']['mergeable_ranks'].items()}
+		dump({
+			'model_name': 'metavoice-1B-v0.1',
+			'stage1': {
+				'tokenizer_info': tokenizer_info
+			},
+			'stage2': {
+				'config': stage2_checkpoint['config'],
+				'meta': stage2_checkpoint['meta'],
+				'model_args': stage2_checkpoint['model_args']
+			}
+		}, f)
+
+class SpeakerEncoder(FAMSpeakerEncoder):
+	def __init__(
+		self,
+		weights_fpath: str,
+		device: Optional[Union[str, torch.device]] = None,
+		verbose: bool = True,
+		eval: bool = False,
+	):
+		nn.Module.__init__(self)
+
+		# Define the network
+		self.lstm = nn.LSTM(mel_n_channels, model_hidden_size, model_num_layers, batch_first=True)
+		self.linear = nn.Linear(model_hidden_size, model_embedding_size)
+		self.relu = nn.ReLU()
+
+		# Get the target device
+		if device is None:
+			device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+		elif isinstance(device, str):
+			device = torch.device(device)
+		self.device = device
+
+		weights_fpath = str(weights_fpath)
+		if weights_fpath.endswith('.safetensors'):
+			load_model(self, weights_fpath)
+		else:
+			checkpoint = torch.load(weights_fpath, map_location='cpu')
+			self.load_state_dict(checkpoint['model_state'], strict=False)
+		self.to(device)
+
+		if eval:
+			self.eval()
+
+def load_safetensors_model(checkpoint_path, spk_emb_ckpt_path, device, precision):
+	##### MODEL
+	with torch.device(device):
+		model = Transformer.from_name('metavoice-1B')
+		load_model(model, checkpoint_path)
+	model = model.to(device=device, dtype=precision)
+
+	###### TOKENIZER
+	with open(f'{os.path.dirname(checkpoint_path)}/config.json', 'r') as f:
+		config = load(f)['stage1']
+	config['tokenizer_info']['mergeable_ranks'] = {b64decode(k): v for k, v in config['tokenizer_info']['mergeable_ranks'].items()}
+	tokenizer_info = config['tokenizer_info']
+	tokenizer = TrainedBPETokeniser(**tokenizer_info)
+
+	###### SPEAKER EMBEDDER
+	smodel = SpeakerEncoder(
+		weights_fpath=spk_emb_ckpt_path,
+		device=device,
+		eval=True,
+		verbose=False,
+	)
+	return model.eval(), tokenizer, smodel
+
+class Model(FAMModel):
+	def _init_model(self):
+		if self.config.init_from == 'safetensors':
+			with open(f'{os.path.dirname(self.config.ckpt_path)}/config.json', 'r') as f:
+				config = load(f)['stage2']
+			self.vocab_sizes = config['model_args']['vocab_sizes']
+			self.checkpoint_config = config['config']
+			config['meta']['tokenizer']['mergeable_ranks'] = {b64decode(k): v for k, v in config['meta']['tokenizer']['mergeable_ranks'].items()}
+
+			self.meta = config['meta']
+			self.load_meta = True
+			self.use_bpe_tokenizer = 'stoi' not in self.meta or 'itos' not in self.meta
+			self.speaker_cond = self.meta.get('speaker_cond')
+
+			speaker_emb_size = None
+			if self.speaker_cond:
+				speaker_emb_size = self.meta['speaker_emb_size']
+
+			model_args = config['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)
+			self.model = GPT(gptconf, speaker_emb_dim=speaker_emb_size)
+			load_model(self.model, self.config.ckpt_path)
+
+		super()._init_model()
+
+class MetaVoiceModel(FAMTTS):
+	def __init__(self, model_name: str, *, seed: int = 1337, output_dir: str = 'outputs', enforce_safetensors: bool = True):
+		self._dtype = get_default_dtype()
+		self._device = get_device()
+
+		if os.path.exists(model_name):
+			if enforce_safetensors:
+				assert all(x in os.listdir(model_name) for x in ('config.json', 'second_stage.safetensors', 'first_stage.safetensors', 'speaker_encoder.safetensors')), 'Model is not compatible with safetensors'
+				self._model_dir = model_name
+			else:
+				print('WARNING: metavoice is allowing the use of non-safetensors models. Ensure you understand the risks of loading untrusted models at https://pytorch.org/docs/stable/generated/torch.load.html')
+				self._model_dir = model_name
+		else:
+			if enforce_safetensors:
+				fs = HfFileSystem()
+				files = [os.path.basename(x) for x in fs.ls(model_name, detail=False)]
+				assert all(x in files for x in ('config.json', 'second_stage.safetensors', 'first_stage.safetensors', 'speaker_encoder.safetensors')), 'Model is not compatible with safetensors'
+				self._model_dir = snapshot_download(repo_id=model_name, allow_patterns='second_stage.safetensors,first_stage.safetensors,speaker_encoder.safetensors,config.json')
+			else:
+				print('WARNING: metavoice is allowing the use of non-safetensors models. Ensure you understand the risks of loading untrusted models at https://pytorch.org/docs/stable/generated/torch.load.html')
+				self._model_dir = snapshot_download(repo_id=model_name)
+
+		self.first_stage_adapter = FlattenedInterleavedEncodec2Codebook(end_of_audio_token=self.END_OF_AUDIO_TOKEN)
+		self.output_dir = output_dir
+		os.makedirs(self.output_dir, exist_ok=True)
+
+		is_safetensors = os.path.exists(f'{self._model_dir}/second_stage.safetensors')
+		second_stage_ckpt_path = f'{self._model_dir}/{"second_stage.safetensors" if is_safetensors else "second_stage.pt"}'
+		config_second_stage = InferenceConfig(
+			ckpt_path=second_stage_ckpt_path,
+			num_samples=1,
+			seed=seed,
+			device=self._device,
+			dtype=self._dtype,
+			compile=False,
+			init_from='safetensors' if is_safetensors else 'resume',
+			output_dir=self.output_dir,
+		)
+		data_adapter_second_stage = TiltedEncodec(end_of_audio_token=self.END_OF_AUDIO_TOKEN)
+		self.llm_second_stage = Model(
+			config_second_stage, TrainedBPETokeniser, EncodecDecoder, data_adapter_fn=data_adapter_second_stage.decode
+		)
+
+		self.enhancer = get_enhancer('df')
+		self.precision = {'float16': torch.float16, 'bfloat16': torch.bfloat16}[self._dtype]
+		build_model_kwargs = {
+			'precision': self.precision,
+			'device': self._device,
+			'compile': False,
+			'compile_prefill': True,
+		}
+		if is_safetensors:
+			fam.llm.fast_inference_utils._load_model = load_safetensors_model
+			checkpoint_path, spk_emb_ckpt_path = Path(f'{self._model_dir}/first_stage.safetensors'), Path(f'{self._model_dir}/speaker_encoder.safetensors')
+		else:
+			checkpoint_path, spk_emb_ckpt_path= Path(f'{self._model_dir}/first_stage.pt'), Path(f'{self._model_dir}/speaker_encoder.pt')
+
+		self.model, self.tokenizer, self.smodel, self.model_size = fam.llm.fast_inference_utils.build_model(
+			checkpoint_path=checkpoint_path,
+			spk_emb_ckpt_path=spk_emb_ckpt_path,
+			**build_model_kwargs
+		)
+
+	@torch.inference_mode()
+	def generate(self, text: str, source: str = 'https://upload.wikimedia.org/wikipedia/commons/e/e1/King_Charles_Addresses_Scottish_Parliament_-_12_September_2022.flac'):
+		self.synthesise(text, source)
+
+	def save(self, path: str):
+		save_model(self.model, os.path.join(path, 'first_stage.safetensors'))
+		save_model(self.smodel, os.path.join(path, 'speaker_encoder.safetensors'))
+		save_model(self.llm_second_stage.model, os.path.join(path, 'second_stage.safetensors'))
+
+	@classmethod
+	def from_hub(cls, path: str):
+		# TODO: TEMPORARY OUTPUT DIR
+		return cls(path, enforce_safetensors=True)

metavoice.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+set -euo pipefail
+
+if ! command -v -- ffmpeg > /dev/null 2>&1; then
+	wget https://johnvansickle.com/ffmpeg/builds/ffmpeg-git-amd64-static.tar.xz
+	wget https://johnvansickle.com/ffmpeg/builds/ffmpeg-git-amd64-static.tar.xz.md5
+	md5sum -c ffmpeg-git-amd64-static.tar.xz.md5
+	tar xf ffmpeg-git-amd64-static.tar.xz
+	[ -z VIRTUAL_ENV ] \
+		&& sudo mv ffmpeg-git-*-static/ffprobe ffmpeg-git-*-static/ffmpeg /usr/local/bin/ \
+		|| mv ffmpeg-git-*-static/ffprobe ffmpeg-git-*-static/ffmpeg "$VIRTUAL_ENV/bin"
+	rm -rf ffmpeg-git-*
+fi
+
+(command -v -- rustup > /dev/null 2>&1) || (wget -qO - https://sh.rustup.rs | sh -s -- -y)
+export PATH="$HOME/.cargo/bin:$PATH"
+
+git clone https://github.com/metavoiceio/metavoice-src.git
+cd metavoice-src
+pip install torch
+pip install -r requirements.txt
+pip install --upgrade torch torchaudio
+pip install .
+cd -
+rm -rf metavoice-src
+
+python -c 'from audiocraft.models import MultiBandDiffusion; MultiBandDiffusion.get_mbd_24khz(bw=6)'