Update inference.py

inference.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import logging
 import json
-from typing import Optional
+from typing import Optional, Dict, Type, Union, List, Tuple
 from pathlib import Path
 from dataclasses import dataclass
 import os
@@ -18,6 +18,8 @@ from mars5.minbpe.codebook import CodebookTokenizer
 from mars5.ar_generate import ar_generate
 from mars5.utils import nuke_weight_norm
 from mars5.trim import trim
+from huggingface_hub import ModelHubMixin, hf_hub_download
+from safetensors import safe_open
 import tempfile
 import logging
 
@@ -64,9 +66,7 @@ class InferenceConfig():
     beam_width: int = 1 # only beam width of 1 is currently supported
     ref_audio_pad: float = 0
 
-
-class Mars5TTS(nn.Module):
-
+class Mars5TTS(nn.Module, ModelHubMixin):
     def __init__(self, ar_ckpt, nar_ckpt, device: str = None) -> None:
         super().__init__()
 
@@ -79,16 +79,16 @@ class Mars5TTS(nn.Module):
 
         # save and load text tokenize
         self.texttok = RegexTokenizer(GPT4_SPLIT_PATTERN)
-        tfn = tempfile.mkstemp(suffix='texttok.model')[1]
-        Path(tfn).write_text(ar_ckpt['vocab']['texttok.model'])
-        self.texttok.load(tfn)
-        os.remove(tfn)
+        texttok_data = io.BytesIO(ar_ckpt['vocab']['texttok.model'].encode('utf-8'))
+        self.texttok.load(texttok_data)
+        texttok_data.close()
+
         # save and load speech tokenizer
-        sfn = tempfile.mkstemp(suffix='speechtok.model')[1]
         self.speechtok = CodebookTokenizer(GPT4_SPLIT_PATTERN)
-        Path(sfn).write_text(ar_ckpt['vocab']['speechtok.model'])
-        self.speechtok.load(sfn)
-        os.remove(sfn)
+        speechtok_data = io.BytesIO(ar_ckpt['vocab']['speechtok.model'].encode('utf-8'))
+        self.speechtok.load(speechtok_data)
+        speechtok_data.close()
+        
         # keep track of tokenization things. 
         self.n_vocab = len(self.texttok.vocab) + len(self.speechtok.vocab)
         self.n_text_vocab = len(self.texttok.vocab) + 1 
@@ -111,7 +111,42 @@ class Mars5TTS(nn.Module):
         self.vocos = Vocos.from_pretrained("charactr/vocos-encodec-24khz").to(self.device).eval()
         nuke_weight_norm(self.codec)
         nuke_weight_norm(self.vocos)
-        
+
+    @classmethod
+    def _from_pretrained(
+        cls: Type["Mars5TTS"],
+        *,
+        model_id: str,
+        revision: Optional[str],
+        cache_dir: Optional[Union[str, Path]],
+        force_download: bool,
+        proxies: Optional[Dict],
+        local_files_only: bool,
+        token: Optional[Union[str, bool]],
+        device: str = None,
+        **model_kwargs,
+    ) -> "Mars5TTS":
+        # Download files from Hub
+        ar_ckpt_path = hf_hub_download(repo_id=model_id, filename="mars5_ar.safetensors", revision=revision, cache_dir=cache_dir, force_download=force_download, proxies=proxies, local_files_only=local_files_only, token=token)
+        nar_ckpt_path = hf_hub_download(repo_id=model_id, filename="mars5_nar.safetensors", revision=revision, cache_dir=cache_dir, force_download=force_download, proxies=proxies, local_files_only=local_files_only, token=token)
+
+        ar_ckpt = {}
+        with safe_open(ar_ckpt_path, framework='pt', device='cpu') as f:
+            metadata = f.metadata()
+            ar_ckpt['vocab'] = {'texttok.model': metadata['texttok.model'], 'speechtok.model': metadata['speechtok.model']}
+            ar_ckpt['model'] = {}
+            for k in f.keys(): ar_ckpt['model'][k] = f.get_tensor(k)
+        nar_ckpt = {}
+        with safe_open(nar_ckpt_path, framework='pt', device='cpu') as f:
+            metadata = f.metadata()
+            nar_ckpt['vocab'] = {'texttok.model': metadata['texttok.model'], 'speechtok.model': metadata['speechtok.model']}
+            nar_ckpt['model'] = {}
+            for k in f.keys(): nar_ckpt['model'][k] = f.get_tensor(k)
+
+
+        # Init
+        return cls(ar_ckpt=ar_ckpt, nar_ckpt=nar_ckpt, device=device)
+
     @torch.inference_mode
     def vocode(self, tokens: Tensor) -> Tensor:
         """ Vocodes tokens of shape (seq_len, n_q) """
@@ -126,6 +161,33 @@ class Mars5TTS(nn.Module):
         wav_diffusion = self.vocos.decode(features, bandwidth_id=bandwidth_id)
         return wav_diffusion.cpu().squeeze()[None]
 
+    @torch.inference_mode
+    def get_speaker_embedding(self, ref_audio: Tensor) -> Tensor:
+        """ Given `ref_audio` (bs, T) audio tensor, compute the implicit speakre embedding of shape (bs, dim). """
+        if ref_audio.dim() == 1: ref_audio = ref_audio[None]
+        spk_reference = self.codec.encode(ref_audio[None].to(self.device))[0][0]
+        spk_reference = spk_reference.permute(0, 2, 1)
+        bs = spk_reference.shape[0]
+        if bs != 1:
+            raise AssertionError(f"Speaker embedding extraction only implemented using for bs=1 currently.")
+        spk_seq = self.codeclm.ref_chunked_emb(spk_reference) # (bs, sl, dim)
+        spk_ref_emb = self.codeclm.spk_identity_emb.weight[None].expand(bs, -1, -1) # (bs, 1, dim)
+
+        spk_seq = torch.cat([spk_ref_emb, spk_seq], dim=1) # (bs, 1+sl, dim)
+        # add pos encoding
+        spk_seq = self.codeclm.pos_embedding(spk_seq)
+        # codebook goes from indices 0->1023, padding is idx 1024 (the 1025th entry)
+        src_key_padding_mask = construct_padding_mask(spk_reference[:, :, 0], 1024) 
+        src_key_padding_mask = torch.cat((
+                                            # append a zero here since we DO want to attend to initial position.
+                                            torch.zeros(src_key_padding_mask.shape[0], 1, dtype=bool, device=src_key_padding_mask.device), 
+                                            src_key_padding_mask
+                                            ), 
+                                            dim=1)
+        # pass through transformer
+        res = self.codeclm.spk_encoder(spk_seq, is_causal=False, src_key_padding_mask=src_key_padding_mask)[:, :1] # select first element -> now (bs, 1, dim).
+        return res.squeeze(1)
+
     @torch.inference_mode
     def tts(self, text: str, ref_audio: Tensor, ref_transcript: Optional[str] = None, 
             cfg: Optional[InferenceConfig] = InferenceConfig()) -> Tensor:
@@ -183,12 +245,12 @@ class Mars5TTS(nn.Module):
         first_codec_idx = prompt.shape[-1] - n_speech_inp + 1
 
         # ---> perform AR code generation
-
         logging.debug(f"Raw acoustic prompt length: {raw_prompt_acoustic_len}")
 
         ar_codes = ar_generate(self.texttok, self.speechtok, self.codeclm, 
                                prompt, spk_ref_codec, first_codec_idx, 
-                               max_len=cfg.generate_max_len_override if cfg.generate_max_len_override > 1 else 2000, 
+                               max_len=cfg.generate_max_len_override if cfg.generate_max_len_override > 1 else 2000,
+                               fp16=True if torch.cuda.is_available() else False,
                                temperature=cfg.temperature, topk=cfg.top_k, top_p=cfg.top_p, typical_p=cfg.typical_p,
                                alpha_frequency=cfg.freq_penalty, alpha_presence=cfg.presence_penalty, penalty_window=cfg.rep_penalty_window,
                                eos_penalty_decay=cfg.eos_penalty_decay, eos_penalty_factor=cfg.eos_penalty_factor,
@@ -211,7 +273,6 @@ class Mars5TTS(nn.Module):
         x_padding_mask = torch.zeros((1, _x.shape[1]), dtype=torch.bool, device=_x.device)
 
         # ---> perform DDPM NAR inference
-
         T = self.default_T
         diff = MultinomialDiffusion(self.diffusion_n_classes, timesteps=T, device=self.device)