| |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| _REPO_DEFAULT = "/scratch/users/ntu/mwong048/FYP-Mamba/external/fish-speech" |
| _REPO_FALLBACK = "/home/max/fish-speech" |
| _CKPT_DEFAULT = "/scratch/users/ntu/mwong048/FYP-Mamba/fish_codec.pth" |
| _CKPT_FALLBACK = "/mnt/c/Users/Max/PycharmProjects/FYP-Mamba/fish_pilot/codec.pth" |
|
|
|
|
| def _resolve(env_var, default, fallback): |
| value = os.environ.get(env_var) |
| if value: |
| return value |
| for candidate in (default, fallback): |
| if Path(candidate).exists(): |
| return candidate |
| return default |
|
|
|
|
| def fish_repo_path(): |
| return _resolve("FISH_SPEECH_REPO", _REPO_DEFAULT, _REPO_FALLBACK) |
|
|
|
|
| def fish_ckpt_path(): |
| return _resolve("FISH_CODEC_PTH", _CKPT_DEFAULT, _CKPT_FALLBACK) |
|
|
|
|
| def load_fish_codec(device): |
| repo = fish_repo_path() |
| if repo not in sys.path: |
| sys.path.insert(0, repo) |
|
|
| from omegaconf import OmegaConf |
| from hydra.utils import instantiate |
|
|
| try: |
| OmegaConf.register_new_resolver("eval", eval) |
| except Exception: |
| pass |
|
|
| cfg = OmegaConf.load(f"{repo}/fish_speech/configs/modded_dac_vq.yaml") |
| model = instantiate(cfg) |
|
|
| ckpt = fish_ckpt_path() |
| sd = torch.load(ckpt, map_location="cpu", weights_only=True, mmap=True) |
| if "state_dict" in sd: |
| sd = sd["state_dict"] |
| if any("generator" in k for k in sd): |
| sd = {k.replace("generator.", ""): v |
| for k, v in sd.items() if "generator." in k} |
| model.load_state_dict(sd, strict=False, assign=True) |
| model.eval().to(device) |
| return model |
|
|
|
|
| def encode_np(model, wav_np, sr, device, max_duration_s=None): |
| import numpy as np |
| import soxr |
|
|
| try: |
| wav = np.asarray(wav_np, dtype=np.float32) |
| if wav.ndim == 2: |
| wav = wav.mean(axis=0) if wav.shape[0] < wav.shape[1] else wav.mean(axis=1) |
| if wav.ndim != 1 or wav.size == 0: |
| return None |
|
|
| if max_duration_s is not None: |
| max_samples = int(max_duration_s * sr) |
| if wav.shape[0] > max_samples: |
| wav = wav[:max_samples] |
|
|
| target_sr = int(model.sample_rate) |
| if int(sr) != target_sr: |
| wav = soxr.resample(wav, int(sr), target_sr) |
| wav = np.ascontiguousarray(wav, dtype=np.float32) |
| if wav.size == 0: |
| return None |
|
|
| x = torch.from_numpy(wav).float()[None, None, :].to(device) |
| lens = torch.tensor([x.shape[-1]], device=device, dtype=torch.long) |
| with torch.no_grad(): |
| indices, idx_lens = model.encode(x, lens) |
| codes = indices[0] |
| if idx_lens is not None: |
| t_valid = int(idx_lens[0]) |
| if 0 < t_valid <= codes.shape[-1]: |
| codes = codes[:, :t_valid] |
| if codes.numel() == 0: |
| return None |
| return [row.tolist() for row in codes.cpu()] |
| except Exception: |
| return None |
|
|
