File size: 3,016 Bytes
6e802a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | # fish_codec_util.py: Loads the Fish S2 codec from a fish-speech checkout (FISH_SPEECH_REPO, FISH_CODEC_PTH).
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
|