Running FishSpeech in a more pythonic way

[saiakarsh-augnito]

I am quite new to fish-speech, and when I went through their documentation, they only provided a CLI version to run their model. I also found it a bit counterintuitive and complex to run those steps. The other issue was when I had to generate multiple audios using CLI, I was loading the model every single time.

If anyone is like me and is more comfortable in the "huggingface way" of running models, I made this script using the existing scripts given in the documentation.

What you get in this script:

• single python script (and easy to understand, phew!)
• model handing is done by hf_hub directly, so we provide either the checkpoint path or model name directly!
• don't need to reload model, rather you can run multiple speech generation in a loop using the same model
• simple SpeechProfile which stores the reference audio codec and uses that to condition speech while generating

How to setup?

• Install fish-speech normally: Installation
• Copy this code (say easy_fish.py and save it in the same folder as fish-speech repo. Otherwise, you need to update the hydra config path accordingly with initialize(version_base="1.3", config_path="fish-speech/fish_speech/configs") (ensure its a relative path)

.
├── female_ns3.wav
├── easy_fish.py
└── fish-speech/

• that's it! since model download is handled by hf_hub, you don't need to worry about that!

How to run?

import soundfile as sf
from easy_fish import FisherTTS

# load fishaudio/openaudio-s1-mini from huggingface
model = FisherTTS('openaudio-s1-mini', half=True)

# no speech profile
wav = model.generate(text='this is an example sentence')

# load a speech profile
spk = model.get_speech_profile('female_ns3.wav')
wav = model.generate(text='this is an example sentence', profile=spk)

# saving the audio
sf.write(f"demo.wav", wav, model.sr)

Hope someone finds this useful! Thank me later :)

easy_fish.py

import os
import time
import torch
import torchaudio
import numpy as np
import soundfile as sf
from loguru import logger
from dataclasses import dataclass
from huggingface_hub import snapshot_download
from typing import Optional

import hydra
from hydra import compose, initialize
from hydra.utils import instantiate
from omegaconf import OmegaConf

from fish_speech.models.text2semantic.inference import init_model, generate_long

########### DAC
# register eval resolver
OmegaConf.register_new_resolver("eval", eval)

def load_model(config_name, checkpoint_path, device="cuda"):
    hydra.core.global_hydra.GlobalHydra.instance().clear()
    # NOTE: using changed relative config path
    with initialize(version_base="1.3", config_path="fish-speech/fish_speech/configs"):
        cfg = compose(config_name=config_name)

    model = instantiate(cfg)
    state_dict = torch.load(
        checkpoint_path, map_location=device, mmap=True, weights_only=True
    )
    if "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]

    if any("generator" in k for k in state_dict):
        state_dict = {
            k.replace("generator.", ""): v
            for k, v in state_dict.items()
            if "generator." in k
        }

    result = model.load_state_dict(state_dict, strict=False, assign=True)
    model.eval()
    model.to(device)

    logger.info(f"Loaded model: {result}")
    return model

# https://speech.fish.audio/inference/#command-line-inference

@dataclass
class SpeechProfile:
    text: str
    code: torch.Tensor

class FisherTTS:
    def __init__(
            self,
            name: str = 'openaudio-s1-mini',
            device: str = 'cuda',
            compile: bool = False,
            seed: int = 42,
            half: bool = False
        ) -> None:
        if not os.path.exists(name):
            repo_id = 'fishaudio/' + name # fishaudio/openaudio-s1-mini
            checkpoint_path = snapshot_download(repo_id)
        else:
            checkpoint_path = name

        precision = torch.half if half else torch.bfloat16
        logger.info("Loading model ...")
        t0 = time.time()
        self.model, self.decode_one_token = init_model(
            checkpoint_path, device, precision, compile=compile
        )
        with torch.device(device):
            self.model.setup_caches(
                max_batch_size=1,
                max_seq_len=self.model.config.max_seq_len,
                dtype=next(self.model.parameters()).dtype,
            )
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.manual_seed(seed)
        logger.info(f"Time to load model: {time.time() - t0:.02f} seconds")

        self.codec_device = 'cpu' # device
        self.codec_model = load_model(config_name='modded_dac_vq', checkpoint_path=checkpoint_path + '/codec.pth', device=self.codec_device)
        self.sr: int = self.codec_model.sample_rate

        self.device = device
        self.compile = compile

    def get_speech_profile(self, reference_audio: str, reference_text: str = '') -> SpeechProfile:
        logger.info(f"Processing in-place reconstruction of {reference_audio}")
        audio, sr = torchaudio.load(str(reference_audio))
        if audio.shape[0] > 1:
            audio = audio.mean(0, keepdim=True)
        audio = torchaudio.functional.resample(audio, sr, self.sr)

        audios = audio[None].to(self.codec_device)
        logger.info(
            f"Loaded audio with {audios.shape[2] / self.sr:.2f} seconds"
        )

        # VQ Encoder
        audio_lengths = torch.tensor([audios.shape[2]], device=self.codec_device, dtype=torch.long)
        indices, indices_lens = self.codec_model.encode(audios, audio_lengths)
        if indices.ndim == 3:
            indices = indices[0]
        logger.info(f"Generated indices of shape {indices.shape}")
        return SpeechProfile(
            text=reference_text,
            code=indices
        )

    def generate(
            self,
            text: str,
            profile: Optional[SpeechProfile] = None,
            num_samples: int = 1,
            max_new_tokens: int = 0,
            top_p: float = 0.8,
            repetition_penalty: float = 1.1,
            temperature: float = 0.8,
            iterative_prompt: bool = True,
            chunk_length: int = 300,
        ) -> np.ndarray:
        prompt_text = None
        prompt_tokens_list = None
        if profile is not None:
            prompt_text = [profile.text]
            prompt_tokens_list = [profile.code]

        with torch.no_grad():
            generator = generate_long(
                model=self.model,
                device=self.device,
                decode_one_token=self.decode_one_token,
                text=text,
                num_samples=num_samples,
                max_new_tokens=max_new_tokens,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                temperature=temperature,
                compile=self.compile,
                iterative_prompt=iterative_prompt,
                chunk_length=chunk_length,
                prompt_text=prompt_text,
                prompt_tokens=prompt_tokens_list,
            )

        codes = []
        for response in generator:
            if response.action == "sample":
                codes.append(response.codes)
                logger.info(f"Sampled text: {response.text}")
            elif response.action == "next":
                logger.info(f"Next sample")
            else:
                logger.error(f"Error: {response}")
        codes = torch.cat(codes, dim=1)

        indices = codes.long().to(device=self.codec_device)
        assert indices.ndim == 2, f"Expected 2D indices, got {indices.ndim}"
        indices_lens = torch.tensor([indices.shape[1]], device=self.codec_device, dtype=torch.long)

        fake_audios: torch.Tensor = self.codec_model.decode(indices, indices_lens)[0] # fake_audios, audio_lengths
        # [1, 1, 71680] -> [71680]
        fake_audio = fake_audios.squeeze().float().detach().cpu().numpy()
        audio_time = fake_audio.shape[0] / self.sr
        logger.info(
            f"Generated audio of shape {fake_audio.shape}, equivalent to {audio_time:.2f} seconds from {indices.shape[1]} features, features/second: {indices.shape[1] / audio_time:.2f}"
        )
        return fake_audio

[keeeers]

Thanks a lot. I'm reading the source code about server part. And I plan to rewrite it with FastAPI to build a stronger self-host server. Your code is pythonic and easy to use. Thanks again!