| | from pathlib import Path
|
| |
|
| | import librosa
|
| | import torch
|
| | import perth
|
| | from huggingface_hub import hf_hub_download
|
| | from safetensors.torch import load_file
|
| |
|
| | from .models.s3tokenizer import S3_SR
|
| | from .models.s3gen import S3GEN_SR, S3Gen
|
| |
|
| |
|
| | REPO_ID = "ResembleAI/chatterbox"
|
| |
|
| |
|
| | class ChatterboxVC:
|
| | ENC_COND_LEN = 6 * S3_SR
|
| | DEC_COND_LEN = 10 * S3GEN_SR
|
| |
|
| | def __init__(
|
| | self,
|
| | s3gen: S3Gen,
|
| | device: str,
|
| | ref_dict: dict=None,
|
| | ):
|
| | self.sr = S3GEN_SR
|
| | self.s3gen = s3gen
|
| | self.device = device
|
| | self.watermarker = perth.PerthImplicitWatermarker()
|
| | if ref_dict is None:
|
| | self.ref_dict = None
|
| | else:
|
| | self.ref_dict = {
|
| | k: v.to(device) if torch.is_tensor(v) else v
|
| | for k, v in ref_dict.items()
|
| | }
|
| |
|
| | @classmethod
|
| | def from_local(cls, ckpt_dir, device) -> 'ChatterboxVC':
|
| | ckpt_dir = Path(ckpt_dir)
|
| |
|
| |
|
| | if device in ["cpu", "mps"]:
|
| | map_location = torch.device('cpu')
|
| | else:
|
| | map_location = None
|
| |
|
| | ref_dict = None
|
| | if (builtin_voice := ckpt_dir / "conds.pt").exists():
|
| | states = torch.load(builtin_voice, map_location=map_location)
|
| | ref_dict = states['gen']
|
| |
|
| | s3gen = S3Gen()
|
| | s3gen.load_state_dict(
|
| | load_file(ckpt_dir / "s3gen.safetensors"), strict=False
|
| | )
|
| | s3gen.to(device).eval()
|
| |
|
| | return cls(s3gen, device, ref_dict=ref_dict)
|
| |
|
| | @classmethod
|
| | def from_pretrained(cls, device) -> 'ChatterboxVC':
|
| |
|
| | if device == "mps" and not torch.backends.mps.is_available():
|
| | if not torch.backends.mps.is_built():
|
| | print("MPS not available because the current PyTorch install was not built with MPS enabled.")
|
| | else:
|
| | print("MPS not available because the current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")
|
| | device = "cpu"
|
| |
|
| | for fpath in ["s3gen.safetensors", "conds.pt"]:
|
| | local_path = hf_hub_download(repo_id=REPO_ID, filename=fpath)
|
| |
|
| | return cls.from_local(Path(local_path).parent, device)
|
| |
|
| | def set_target_voice(self, wav_fpath):
|
| |
|
| | s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
|
| |
|
| | s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
|
| | self.ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
|
| |
|
| | def generate(
|
| | self,
|
| | audio,
|
| | target_voice_path=None,
|
| | ):
|
| | if target_voice_path:
|
| | self.set_target_voice(target_voice_path)
|
| | else:
|
| | assert self.ref_dict is not None, "Please `prepare_conditionals` first or specify `target_voice_path`"
|
| |
|
| | with torch.inference_mode():
|
| | audio_16, _ = librosa.load(audio, sr=S3_SR)
|
| | audio_16 = torch.from_numpy(audio_16).float().to(self.device)[None, ]
|
| |
|
| | s3_tokens, _ = self.s3gen.tokenizer(audio_16)
|
| | wav, _ = self.s3gen.inference(
|
| | speech_tokens=s3_tokens,
|
| | ref_dict=self.ref_dict,
|
| | )
|
| | wav = wav.squeeze(0).detach().cpu().numpy()
|
| | watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
|
| | return torch.from_numpy(watermarked_wav).unsqueeze(0) |
