xinference 1.1.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

xinference/thirdparty/fish_speech/fish_speech/train.py

@@ -6,7 +6,7 @@ from typing import Optional
 
 import hydra
 import lightning as L
-# import pyrootutils
+import pyrootutils
 import torch
 from lightning import Callback, LightningDataModule, LightningModule, Trainer
 from lightning.pytorch.loggers import Logger
@@ -18,7 +18,7 @@ os.environ.pop("SLURM_JOB_NAME", None)
 os.environ.pop("SLURM_NTASKS_PER_NODE", None)
 
 # register eval resolver and root
-# pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
 
 # Allow TF32 on Ampere GPUs
 torch.set_float32_matmul_precision("high")

xinference/thirdparty/fish_speech/fish_speech/webui/manage.py

@@ -176,7 +176,7 @@ def change_infer(
         p_infer = subprocess.Popen(
             [
                 PYTHON,
-                "tools/webui.py",
+                "tools/run_webui.py",
                 "--decoder-checkpoint-path",
                 infer_decoder_model,
                 "--decoder-config-name",

xinference/thirdparty/fish_speech/tools/{post_api.py → api_client.py}

@@ -69,10 +69,6 @@ def parse_args():
     parser.add_argument(
         "--format", type=str, choices=["wav", "mp3", "flac"], default="wav"
     )
-    parser.add_argument(
-        "--mp3_bitrate", type=int, choices=[64, 128, 192], default=64, help="kHz"
-    )
-    parser.add_argument("--opus_bitrate", type=int, default=-1000)
     parser.add_argument(
         "--latency",
         type=str,
@@ -83,7 +79,7 @@ def parse_args():
     parser.add_argument(
         "--max_new_tokens",
         type=int,
-        default=0,
+        default=1024,
         help="Maximum new tokens to generate. \n0 means no limit.",
     )
     parser.add_argument(
@@ -112,11 +108,9 @@ def parse_args():
     parser.add_argument(
         "--use_memory_cache",
         type=str,
-        default="never",
-        choices=["on-demand", "never"],
-        help="Cache encoded references codes in memory.\n"
-        "If `on-demand`, the server will use cached encodings\n "
-        "instead of encoding reference audio again.",
+        default="off",
+        choices=["on", "off"],
+        help="Cache encoded references codes in memory.\n",
     )
     parser.add_argument(
         "--seed",
@@ -154,14 +148,14 @@ if __name__ == "__main__":
     data = {
         "text": args.text,
         "references": [
-            ServeReferenceAudio(audio=ref_audio, text=ref_text)
+            ServeReferenceAudio(
+                audio=ref_audio if ref_audio is not None else b"", text=ref_text
+            )
             for ref_text, ref_audio in zip(ref_texts, byte_audios)
         ],
         "reference_id": idstr,
         "normalize": args.normalize,
         "format": args.format,
-        "mp3_bitrate": args.mp3_bitrate,
-        "opus_bitrate": args.opus_bitrate,
         "max_new_tokens": args.max_new_tokens,
         "chunk_length": args.chunk_length,
         "top_p": args.top_p,

xinference/thirdparty/fish_speech/tools/api_server.py

@@ -0,0 +1,98 @@
+from threading import Lock
+
+import pyrootutils
+import uvicorn
+from kui.asgi import FactoryClass, HTTPException, HttpRoute, Kui, OpenAPI, Routes
+from loguru import logger
+
+pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from tools.server.api_utils import MsgPackRequest, parse_args
+from tools.server.exception_handler import ExceptionHandler
+from tools.server.model_manager import ModelManager
+from tools.server.views import (
+    ASRView,
+    ChatView,
+    HealthView,
+    TTSView,
+    VQGANDecodeView,
+    VQGANEncodeView,
+)
+
+
+class API(ExceptionHandler):
+    def __init__(self):
+        self.args = parse_args()
+        self.routes = [
+            ("/v1/health", HealthView),
+            ("/v1/vqgan/encode", VQGANEncodeView),
+            ("/v1/vqgan/decode", VQGANDecodeView),
+            ("/v1/asr", ASRView),
+            ("/v1/tts", TTSView),
+            ("/v1/chat", ChatView),
+        ]
+        self.routes = Routes([HttpRoute(path, view) for path, view in self.routes])
+
+        self.openapi = OpenAPI(
+            {
+                "title": "Fish Speech API",
+                "version": "1.5.0",
+            },
+        ).routes
+
+        # Initialize the app
+        self.app = Kui(
+            routes=self.routes + self.openapi[1:],  # Remove the default route
+            exception_handlers={
+                HTTPException: self.http_exception_handler,
+                Exception: self.other_exception_handler,
+            },
+            factory_class=FactoryClass(http=MsgPackRequest),
+            cors_config={},
+        )
+
+        # Add the state variables
+        self.app.state.lock = Lock()
+        self.app.state.device = self.args.device
+        self.app.state.max_text_length = self.args.max_text_length
+
+        # Associate the app with the model manager
+        self.app.on_startup(self.initialize_app)
+
+    async def initialize_app(self, app: Kui):
+        # Make the ModelManager available to the views
+        app.state.model_manager = ModelManager(
+            mode=self.args.mode,
+            device=self.args.device,
+            half=self.args.half,
+            compile=self.args.compile,
+            asr_enabled=self.args.load_asr_model,
+            llama_checkpoint_path=self.args.llama_checkpoint_path,
+            decoder_checkpoint_path=self.args.decoder_checkpoint_path,
+            decoder_config_name=self.args.decoder_config_name,
+        )
+
+        logger.info(f"Startup done, listening server at http://{self.args.listen}")
+
+
+# Each worker process created by Uvicorn has its own memory space,
+# meaning that models and variables are not shared between processes.
+# Therefore, any variables (like `llama_queue` or `decoder_model`)
+# will not be shared across workers.
+
+# Multi-threading for deep learning can cause issues, such as inconsistent
+# outputs if multiple threads access the same buffers simultaneously.
+# Instead, it's better to use multiprocessing or independent models per thread.
+
+if __name__ == "__main__":
+
+    api = API()
+    host, port = api.args.listen.split(":")
+
+    uvicorn.run(
+        api.app,
+        host=host,
+        port=int(port),
+        workers=api.args.workers,
+        log_level="info",
+    )

xinference/thirdparty/fish_speech/tools/download_models.py

@@ -22,14 +22,14 @@ def check_and_download_files(repo_id, file_list, local_dir):
 
 
 # 1st
-repo_id_1 = "fishaudio/fish-speech-1.4"
-local_dir_1 = "./checkpoints/fish-speech-1.4"
+repo_id_1 = "fishaudio/fish-speech-1.5"
+local_dir_1 = "./checkpoints/fish-speech-1.5"
 files_1 = [
+    "gitattributes",
     "model.pth",
     "README.md",
-    "special_tokens_map.json",
-    "tokenizer_config.json",
-    "tokenizer.json",
+    "special_tokens.json",
+    "tokenizer.tiktoken",
     "config.json",
     "firefly-gan-vq-fsq-8x1024-21hz-generator.pth",
 ]

xinference/thirdparty/fish_speech/tools/fish_e2e.py

@@ -14,8 +14,8 @@ import ormsgpack
 import soundfile as sf
 
 from .schema import (
+    ServeChatRequest,
     ServeMessage,
-    ServeRequest,
     ServeTextPart,
     ServeVQGANDecodeRequest,
     ServeVQGANEncodeRequest,
@@ -163,7 +163,7 @@ class FishE2EAgent:
         else:
             user_codes = None
 
-        request = ServeRequest(
+        request = ServeChatRequest(
             messages=prev_messages
             + (
                 [

xinference/thirdparty/fish_speech/tools/inference_engine/__init__.py

@@ -0,0 +1,192 @@
+import gc
+import queue
+from typing import Generator
+
+import numpy as np
+import torch
+from loguru import logger
+
+from fish_speech.models.vqgan.modules.firefly import FireflyArchitecture
+from fish_speech.text.chn_text_norm.text import Text as ChnNormedText
+from fish_speech.utils import autocast_exclude_mps, set_seed
+from tools.inference_engine.reference_loader import ReferenceLoader
+from tools.inference_engine.utils import InferenceResult, wav_chunk_header
+from tools.inference_engine.vq_manager import VQManager
+from tools.llama.generate import (
+    GenerateRequest,
+    GenerateResponse,
+    WrappedGenerateResponse,
+)
+from tools.schema import ServeTTSRequest
+
+
+class TTSInferenceEngine(ReferenceLoader, VQManager):
+
+    def __init__(
+        self,
+        llama_queue: queue.Queue,
+        decoder_model: FireflyArchitecture,
+        precision: torch.dtype,
+        compile: bool,
+    ) -> None:
+
+        super().__init__()
+
+        self.llama_queue = llama_queue
+        self.decoder_model = decoder_model
+        self.precision = precision
+        self.compile = compile
+
+    @torch.inference_mode()
+    def inference(self, req: ServeTTSRequest) -> Generator[InferenceResult, None, None]:
+        """
+        Main inference function:
+        - Loads the reference audio and text.
+        - Calls the LLAMA model for inference.
+        - Decodes the VQ tokens to audio.
+        """
+
+        ref_id: str | None = req.reference_id
+        prompt_tokens, prompt_texts = [], []
+        # Load the reference audio and text based on id or hash
+        if ref_id is not None:
+            prompt_tokens, prompt_texts = self.load_by_id(ref_id, req.use_memory_cache)
+
+        elif req.references:
+            prompt_tokens, prompt_texts = self.load_by_hash(
+                req.references, req.use_memory_cache
+            )
+
+        # Set the random seed if provided
+        if req.seed is not None:
+            set_seed(req.seed)
+            logger.warning(f"set seed: {req.seed}")
+
+        # Get the symbolic tokens from the LLAMA model
+        response_queue = self.send_Llama_request(req, prompt_tokens, prompt_texts)
+
+        # Get the sample rate from the decoder model
+        sample_rate = self.decoder_model.spec_transform.sample_rate
+
+        # If streaming, send the header
+        # if req.streaming:
+        #     yield InferenceResult(
+        #         code="header",
+        #         audio=(sample_rate, wav_chunk_header(sample_rate=sample_rate)),
+        #         error=None,
+        #     )
+
+        segments = []
+
+        while True:
+            # Get the response from the LLAMA model
+            wrapped_result: WrappedGenerateResponse = response_queue.get()
+            if wrapped_result.status == "error":
+                yield InferenceResult(
+                    code="error",
+                    audio=None,
+                    error=(
+                        wrapped_result.response
+                        if isinstance(wrapped_result.response, Exception)
+                        else Exception("Unknown error")
+                    ),
+                )
+                break
+
+            # Check the response type
+            if not isinstance(wrapped_result.response, GenerateResponse):
+                raise TypeError(
+                    "Expected GenerateResponse, got {type(wrapped_result.response).__name__}"
+                )
+
+            result: GenerateResponse = wrapped_result.response
+            if result.action != "next":
+                segment = self.get_audio_segment(result)
+
+                if req.streaming:  # Used only by the API server
+                    yield InferenceResult(
+                        code="segment",
+                        audio=(sample_rate, segment),
+                        error=None,
+                    )
+                segments.append(segment)
+            else:
+                break
+
+        # Clean up the memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            gc.collect()
+
+        # Edge case: no audio generated
+        if len(segments) == 0:
+            yield InferenceResult(
+                code="error",
+                audio=None,
+                error=RuntimeError("No audio generated, please check the input text."),
+            )
+        else:
+            # Streaming or not, return the final audio
+            audio = np.concatenate(segments, axis=0)
+            yield InferenceResult(
+                code="final",
+                audio=(sample_rate, audio),
+                error=None,
+            )
+
+        return None
+
+    def send_Llama_request(
+        self, req: ServeTTSRequest, prompt_tokens: list, prompt_texts: list
+    ) -> queue.Queue:
+        """
+        Send a request to the LLAMA model to generate the symbolic tokens.
+        """
+
+        # Prepare the request
+        request = dict(
+            device=self.decoder_model.device,
+            max_new_tokens=req.max_new_tokens,
+            text=(
+                req.text
+                if not req.normalize
+                else ChnNormedText(raw_text=req.text).normalize()
+            ),
+            top_p=req.top_p,
+            repetition_penalty=req.repetition_penalty,
+            temperature=req.temperature,
+            compile=self.compile,
+            iterative_prompt=req.chunk_length > 0,
+            chunk_length=req.chunk_length,
+            max_length=4096,
+            prompt_tokens=prompt_tokens,
+            prompt_text=prompt_texts,
+        )
+
+        # Create a queue to get the response
+        response_queue = queue.Queue()
+
+        # Send the request to the LLAMA model
+        self.llama_queue.put(
+            GenerateRequest(
+                request=request,
+                response_queue=response_queue,
+            )
+        )
+
+        return response_queue
+
+    def get_audio_segment(self, result: GenerateResponse) -> np.ndarray:
+        """
+        Decode the VQ tokens to audio.
+        """
+
+        # Don't use autocast on MPS devices
+        with autocast_exclude_mps(
+            device_type=self.decoder_model.device.type, dtype=self.precision
+        ):
+            # Decode the symbolic tokens to audio
+            segment = self.decode_vq_tokens(codes=result.codes)
+
+        # Convert the audio to numpy
+        return segment.float().cpu().numpy()

xinference/thirdparty/fish_speech/tools/inference_engine/reference_loader.py

@@ -0,0 +1,125 @@
+import io
+from hashlib import sha256
+from pathlib import Path
+from typing import Callable, Literal, Tuple
+
+import torch
+import torchaudio
+from loguru import logger
+
+from fish_speech.models.vqgan.modules.firefly import FireflyArchitecture
+from tools.file import AUDIO_EXTENSIONS, audio_to_bytes, list_files, read_ref_text
+from tools.schema import ServeReferenceAudio
+
+
+class ReferenceLoader:
+
+    def __init__(self) -> None:
+        """
+        Component of the TTSInferenceEngine class.
+        Loads and manages the cache for the reference audio and text.
+        """
+        self.ref_by_id: dict = {}
+        self.ref_by_hash: dict = {}
+
+        # Make Pylance happy (attribut/method not defined...)
+        self.decoder_model: FireflyArchitecture
+        self.encode_reference: Callable
+
+        # Define the torchaudio backend
+        backends = torchaudio.list_audio_backends()
+        if "ffmpeg" in backends:
+            self.backend = "ffmpeg"
+        else:
+            self.backend = "soundfile"
+
+    def load_by_id(
+        self,
+        id: str,
+        use_cache: Literal["on", "off"],
+    ) -> Tuple:
+
+        # Load the references audio and text by id
+        ref_folder = Path("references") / id
+        ref_folder.mkdir(parents=True, exist_ok=True)
+        ref_audios = list_files(
+            ref_folder, AUDIO_EXTENSIONS, recursive=True, sort=False
+        )
+
+        if use_cache == "off" or id not in self.ref_by_id:
+            # If the references are not already loaded, encode them
+            prompt_tokens = [
+                self.encode_reference(
+                    # decoder_model=self.decoder_model,
+                    reference_audio=audio_to_bytes(str(ref_audio)),
+                    enable_reference_audio=True,
+                )
+                for ref_audio in ref_audios
+            ]
+            prompt_texts = [
+                read_ref_text(str(ref_audio.with_suffix(".lab")))
+                for ref_audio in ref_audios
+            ]
+            self.ref_by_id[id] = (prompt_tokens, prompt_texts)
+
+        else:
+            # Reuse already encoded references
+            logger.info("Use same references")
+            prompt_tokens, prompt_texts = self.ref_by_id[id]
+
+        return prompt_tokens, prompt_texts
+
+    def load_by_hash(
+        self,
+        references: list[ServeReferenceAudio],
+        use_cache: Literal["on", "off"],
+    ) -> Tuple:
+
+        # Load the references audio and text by hash
+        audio_hashes = [sha256(ref.audio).hexdigest() for ref in references]
+
+        cache_used = False
+        prompt_tokens, prompt_texts = [], []
+        for i, ref in enumerate(references):
+            if use_cache == "off" or audio_hashes[i] not in self.ref_by_hash:
+                # If the references are not already loaded, encode them
+                prompt_tokens.append(
+                    self.encode_reference(
+                        reference_audio=ref.audio,
+                        enable_reference_audio=True,
+                    )
+                )
+                prompt_texts.append(ref.text)
+                self.ref_by_hash[audio_hashes[i]] = (prompt_tokens, prompt_texts)
+
+            else:
+                # Reuse already encoded references
+                prompt_tokens, prompt_texts = self.ref_by_hash[audio_hashes[i]]
+                cache_used = True
+
+        if cache_used:
+            logger.info("Use same references")
+
+        return prompt_tokens, prompt_texts
+
+    def load_audio(self, reference_audio, sr):
+        """
+        Load the audio data from a file or bytes.
+        """
+        if len(reference_audio) > 255 or not Path(reference_audio).exists():
+            audio_data = reference_audio
+            reference_audio = io.BytesIO(audio_data)
+
+        waveform, original_sr = torchaudio.load(reference_audio, backend=self.backend)
+
+        if waveform.shape[0] > 1:
+            waveform = torch.mean(waveform, dim=0, keepdim=True)
+
+        if original_sr != sr:
+            resampler = torchaudio.transforms.Resample(
+                orig_freq=original_sr, new_freq=sr
+            )
+            waveform = resampler(waveform)
+
+        audio = waveform.squeeze().numpy()
+        return audio

xinference/thirdparty/fish_speech/tools/inference_engine/utils.py

@@ -0,0 +1,39 @@
+import io
+import wave
+from dataclasses import dataclass
+from typing import Literal, Optional, Tuple
+
+import numpy as np
+
+from fish_speech.text.chn_text_norm.text import Text as ChnNormedText
+
+
+@dataclass
+class InferenceResult:
+    code: Literal["header", "segment", "error", "final"]
+    audio: Optional[Tuple[int, np.ndarray | bytes]]
+    error: Optional[Exception]
+
+
+def normalize_text(user_input: str, use_normalization: bool) -> str:
+    """Normalize user input text if needed."""
+    if use_normalization:
+        return ChnNormedText(raw_text=user_input).normalize()
+    else:
+        return user_input
+
+
+def wav_chunk_header(
+    sample_rate: int = 44100, bit_depth: int = 16, channels: int = 1
+) -> bytes:
+    buffer = io.BytesIO()
+
+    with wave.open(buffer, "wb") as wav_file:
+        wav_file.setnchannels(channels)
+        wav_file.setsampwidth(bit_depth // 8)
+        wav_file.setframerate(sample_rate)
+
+    wav_header_bytes = buffer.getvalue()
+    buffer.close()
+
+    return wav_header_bytes

xinference/thirdparty/fish_speech/tools/inference_engine/vq_manager.py

@@ -0,0 +1,57 @@
+from typing import Callable
+
+import torch
+from loguru import logger
+
+from fish_speech.models.vqgan.modules.firefly import FireflyArchitecture
+
+
+class VQManager:
+
+    def __init__(self):
+        # Make Pylance happy (attribut/method not defined...)
+        self.decoder_model: FireflyArchitecture
+        self.load_audio: Callable
+
+    def decode_vq_tokens(self, codes):
+        feature_lengths = torch.tensor(
+            [codes.shape[1]], device=self.decoder_model.device
+        )
+        logger.info(f"VQ features: {codes.shape}")
+
+        if isinstance(self.decoder_model, FireflyArchitecture):
+            return self.decoder_model.decode(
+                indices=codes[None],
+                feature_lengths=feature_lengths,
+            )[0].squeeze()
+
+        raise ValueError(f"Unknown model type: {type(self.decoder_model)}")
+
+    def encode_reference(self, reference_audio, enable_reference_audio):
+        if enable_reference_audio and reference_audio is not None:
+            # Load audios, and prepare basic info here
+            reference_audio_content = self.load_audio(
+                reference_audio, self.decoder_model.spec_transform.sample_rate
+            )
+
+            audios = torch.from_numpy(reference_audio_content).to(
+                self.decoder_model.device
+            )[None, None, :]
+            audio_lengths = torch.tensor(
+                [audios.shape[2]], device=self.decoder_model.device, dtype=torch.long
+            )
+            logger.info(
+                f"Loaded audio with {audios.shape[2] / self.decoder_model.spec_transform.sample_rate:.2f} seconds"
+            )
+
+            # VQ Encoder
+            if isinstance(self.decoder_model, FireflyArchitecture):
+                prompt_tokens = self.decoder_model.encode(audios, audio_lengths)[0][0]
+                logger.info(f"Encoded prompt: {prompt_tokens.shape}")
+            else:
+                raise ValueError(f"Unknown model type: {type(self.decoder_model)}")
+        else:
+            prompt_tokens = None
+            logger.info("No reference audio provided")
+
+        return prompt_tokens

xinference/thirdparty/fish_speech/tools/llama/eval_in_context.py

@@ -1,11 +1,11 @@
-# import pyrootutils
+import pyrootutils
 import torch
 import torch.nn.functional as F
 from matplotlib import pyplot as plt
 from transformers import AutoTokenizer
 
 # register eval resolver and root
-# pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
 
 from torch.utils.data import DataLoader