xinference 0.16.3__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/thirdparty/fish_speech/tools/server/api_utils.py

@@ -0,0 +1,75 @@
+from argparse import ArgumentParser
+from http import HTTPStatus
+from typing import Annotated, Any
+
+import ormsgpack
+from baize.datastructures import ContentType
+from kui.asgi import HTTPException, HttpRequest
+
+from tools.inference_engine import TTSInferenceEngine
+from tools.schema import ServeTTSRequest
+from tools.server.inference import inference_wrapper as inference
+
+
+def parse_args():
+    parser = ArgumentParser()
+    parser.add_argument("--mode", type=str, choices=["agent", "tts"], default="tts")
+    parser.add_argument("--load-asr-model", action="store_true")
+    parser.add_argument(
+        "--llama-checkpoint-path",
+        type=str,
+        default="checkpoints/fish-speech-1.5",
+    )
+    parser.add_argument(
+        "--decoder-checkpoint-path",
+        type=str,
+        default="checkpoints/fish-speech-1.5/firefly-gan-vq-fsq-8x1024-21hz-generator.pth",
+    )
+    parser.add_argument("--decoder-config-name", type=str, default="firefly_gan_vq")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--half", action="store_true")
+    parser.add_argument("--compile", action="store_true")
+    parser.add_argument("--max-text-length", type=int, default=0)
+    parser.add_argument("--listen", type=str, default="127.0.0.1:8080")
+    parser.add_argument("--workers", type=int, default=1)
+
+    return parser.parse_args()
+
+
+class MsgPackRequest(HttpRequest):
+    async def data(
+        self,
+    ) -> Annotated[
+        Any, ContentType("application/msgpack"), ContentType("application/json")
+    ]:
+        if self.content_type == "application/msgpack":
+            return ormsgpack.unpackb(await self.body)
+
+        elif self.content_type == "application/json":
+            return await self.json
+
+        raise HTTPException(
+            HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
+            headers={"Accept": "application/msgpack, application/json"},
+        )
+
+
+async def inference_async(req: ServeTTSRequest, engine: TTSInferenceEngine):
+    for chunk in inference(req, engine):
+        if isinstance(chunk, bytes):
+            yield chunk
+
+
+async def buffer_to_async_generator(buffer):
+    yield buffer
+
+
+def get_content_type(audio_format):
+    if audio_format == "wav":
+        return "audio/wav"
+    elif audio_format == "flac":
+        return "audio/flac"
+    elif audio_format == "mp3":
+        return "audio/mpeg"
+    else:
+        return "application/octet-stream"

xinference/thirdparty/fish_speech/tools/server/exception_handler.py

@@ -0,0 +1,27 @@
+import traceback
+from http import HTTPStatus
+
+from kui.asgi import HTTPException, JSONResponse
+
+
+class ExceptionHandler:
+
+    async def http_exception_handler(self, exc: HTTPException):
+        return JSONResponse(
+            dict(
+                statusCode=exc.status_code,
+                message=exc.content,
+                error=HTTPStatus(exc.status_code).phrase,
+            ),
+            exc.status_code,
+            exc.headers,
+        )
+
+    async def other_exception_handler(self, exc: Exception):
+        traceback.print_exc()
+
+        status = HTTPStatus.INTERNAL_SERVER_ERROR
+        return JSONResponse(
+            dict(statusCode=status, message=str(exc), error=status.phrase),
+            status,
+        )

xinference/thirdparty/fish_speech/tools/server/inference.py

@@ -0,0 +1,45 @@
+from http import HTTPStatus
+
+import numpy as np
+from kui.asgi import HTTPException
+
+from tools.inference_engine import TTSInferenceEngine
+from tools.schema import ServeTTSRequest
+
+AMPLITUDE = 32768  # Needs an explaination
+
+
+def inference_wrapper(req: ServeTTSRequest, engine: TTSInferenceEngine):
+    """
+    Wrapper for the inference function.
+    Used in the API server.
+    """
+    count = 0
+    for result in engine.inference(req):
+        match result.code:
+            case "header":
+                if isinstance(result.audio, tuple):
+                    yield result.audio[1]
+
+            case "error":
+                raise HTTPException(
+                    HTTPStatus.INTERNAL_SERVER_ERROR,
+                    content=str(result.error),
+                )
+
+            case "segment":
+                count += 1
+                if isinstance(result.audio, tuple):
+                    yield (result.audio[1] * AMPLITUDE).astype(np.int16).tobytes()
+
+            case "final":
+                count += 1
+                if isinstance(result.audio, tuple):
+                    yield result.audio[1]
+                return None  # Stop the generator
+
+    if count == 0:
+        raise HTTPException(
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            content="No audio generated, please check the input text.",
+        )

xinference/thirdparty/fish_speech/tools/server/model_manager.py

@@ -0,0 +1,122 @@
+import torch
+from funasr import AutoModel
+from loguru import logger
+
+from tools.inference_engine import TTSInferenceEngine
+from tools.llama.generate import (
+    launch_thread_safe_queue,
+    launch_thread_safe_queue_agent,
+)
+from tools.schema import ServeTTSRequest
+from tools.server.inference import inference_wrapper as inference
+from tools.vqgan.inference import load_model as load_decoder_model
+
+ASR_MODEL_NAME = "iic/SenseVoiceSmall"
+
+
+class ModelManager:
+    def __init__(
+        self,
+        mode: str,
+        device: str,
+        half: bool,
+        compile: bool,
+        asr_enabled: bool,
+        llama_checkpoint_path: str,
+        decoder_checkpoint_path: str,
+        decoder_config_name: str,
+    ) -> None:
+
+        self.mode = mode
+        self.device = device
+        self.half = half
+        self.compile = compile
+
+        self.precision = torch.half if half else torch.bfloat16
+
+        # Check if MPS or CUDA is available
+        if torch.backends.mps.is_available():
+            self.device = "mps"
+            logger.info("mps is available, running on mps.")
+        elif not torch.cuda.is_available():
+            self.device = "cpu"
+            logger.info("CUDA is not available, running on CPU.")
+
+        # Load the ASR model if enabled
+        if asr_enabled:
+            self.load_asr_model(self.device)
+
+        # Load the TTS models
+        self.load_llama_model(
+            llama_checkpoint_path, self.device, self.precision, self.compile, self.mode
+        )
+        self.load_decoder_model(
+            decoder_config_name, decoder_checkpoint_path, self.device
+        )
+        self.tts_inference_engine = TTSInferenceEngine(
+            llama_queue=self.llama_queue,
+            decoder_model=self.decoder_model,
+            precision=self.precision,
+            compile=self.compile,
+        )
+
+        # Warm up the models
+        if self.mode == "tts":
+            self.warm_up(self.tts_inference_engine)
+
+    def load_asr_model(self, device, hub="ms") -> None:
+        self.asr_model = AutoModel(
+            model=ASR_MODEL_NAME,
+            device=device,
+            disable_pbar=True,
+            hub=hub,
+        )
+        logger.info("ASR model loaded.")
+
+    def load_llama_model(
+        self, checkpoint_path, device, precision, compile, mode
+    ) -> None:
+
+        if mode == "tts":
+            self.llama_queue = launch_thread_safe_queue(
+                checkpoint_path=checkpoint_path,
+                device=device,
+                precision=precision,
+                compile=compile,
+            )
+        elif mode == "agent":
+            self.llama_queue, self.tokenizer, self.config = (
+                launch_thread_safe_queue_agent(
+                    checkpoint_path=checkpoint_path,
+                    device=device,
+                    precision=precision,
+                    compile=compile,
+                )
+            )
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+
+        logger.info("LLAMA model loaded.")
+
+    def load_decoder_model(self, config_name, checkpoint_path, device) -> None:
+        self.decoder_model = load_decoder_model(
+            config_name=config_name,
+            checkpoint_path=checkpoint_path,
+            device=device,
+        )
+        logger.info("Decoder model loaded.")
+
+    def warm_up(self, tts_inference_engine) -> None:
+        request = ServeTTSRequest(
+            text="Hello world.",
+            references=[],
+            reference_id=None,
+            max_new_tokens=1024,
+            chunk_length=200,
+            top_p=0.7,
+            repetition_penalty=1.2,
+            temperature=0.7,
+            format="wav",
+        )
+        list(inference(request, tts_inference_engine))
+        logger.info("Models warmed up.")

xinference/thirdparty/fish_speech/tools/server/model_utils.py

@@ -0,0 +1,129 @@
+import io
+import re
+
+import librosa
+import torch
+import torchaudio
+from cachetools import LRUCache, cached
+
+CACHE_MAXSIZE = 10000
+MICRO_BATCH_SIZE = 8
+ASR_SAMPLE_RATE = 16000
+HUGE_GAP_THRESHOLD = 4000
+
+
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.half)
+def batch_encode(model, audios_list: list[bytes]):
+    audios: list[torch.Tensor] = [
+        (
+            torch.from_numpy(
+                librosa.load(io.BytesIO(audio), sr=model.spec_transform.sample_rate)[0]
+            )[None]
+            if isinstance(audio, bytes)
+            else audio
+        )
+        for audio in audios_list
+    ]
+
+    lengths = torch.tensor([audio.shape[-1] for audio in audios], device=model.device)
+    max_length = lengths.max().item()
+
+    print(f"Encode max length: {max_length / model.spec_transform.sample_rate:.2f}s")
+
+    padded = torch.stack(
+        [
+            torch.nn.functional.pad(audio, (0, int(max_length - audio.shape[-1])))
+            for audio in audios
+        ]
+    ).to(model.device)
+
+    features, feature_lengths = model.encode(padded, audio_lengths=lengths)
+    features, feature_lengths = features.cpu(), feature_lengths.cpu()
+
+    return [feature[..., :length] for feature, length in zip(features, feature_lengths)]
+
+
+@cached(
+    cache=LRUCache(maxsize=CACHE_MAXSIZE),
+    key=lambda model, audios: (model.device, tuple(audios)),
+)
+def cached_vqgan_batch_encode(model, audios: list[bytes]):
+    return batch_encode(model, audios)
+
+
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.half)
+def vqgan_decode(model, features):
+    lengths = torch.tensor(
+        [feature.shape[-1] for feature in features], device=model.device
+    )
+    max_length = lengths.max().item()
+    padded = torch.stack(
+        [
+            torch.nn.functional.pad(feature, (0, max_length - feature.shape[-1]))
+            for feature in features
+        ]
+    ).to(model.device)
+
+    # If bs too large, we do micro batch decode
+    audios, audio_lengths = [], []
+    for i in range(0, padded.shape[0], MICRO_BATCH_SIZE):
+        audio, audio_length = model.decode(
+            padded[i : i + MICRO_BATCH_SIZE],
+            feature_lengths=lengths[i : i + MICRO_BATCH_SIZE],
+        )
+        audios.append(audio)
+        audio_lengths.append(audio_length)
+    audios = torch.cat(audios, dim=0)
+    audio_lengths = torch.cat(audio_lengths, dim=0)
+    audios, audio_lengths = audios.cpu(), audio_lengths.cpu()
+
+    return [audio[..., :length].numpy() for audio, length in zip(audios, audio_lengths)]
+
+
+@torch.no_grad()
+def batch_asr(model, lock, audios, sr, language="auto"):
+    resampled_audios = []
+    for audio in audios:
+        audio = torchaudio.functional.resample(audio, sr, ASR_SAMPLE_RATE)
+        assert audio.ndim == 1
+        resampled_audios.append(audio)
+
+    with lock:
+        res = model.generate(
+            input=resampled_audios,
+            batch_size=len(resampled_audios),
+            language=language,
+            use_itn=True,
+        )
+
+    results = []
+    for r, audio in zip(res, audios):
+        text = r["text"]
+        text = re.sub(r"<\|.*?\|>", "", text)
+        duration = len(audio) / sr * 1000
+        huge_gap = False
+
+        if "timestamp" in r and len(r["timestamp"]) > 2:
+            for timestamp_a, timestamp_b in zip(
+                r["timestamp"][:-1], r["timestamp"][1:]
+            ):
+                # If there is a gap of more than 4 seconds, we consider it as a huge gap
+                if timestamp_b[0] - timestamp_a[1] > HUGE_GAP_THRESHOLD:
+                    huge_gap = True
+                    break
+
+            # Doesn't make sense to have a huge gap at the end
+            if duration - r["timestamp"][-1][1] > HUGE_GAP_THRESHOLD:
+                huge_gap = True
+
+        results.append(
+            {
+                "text": text,
+                "duration": duration,
+                "huge_gap": huge_gap,
+            }
+        )
+
+    return results

xinference/thirdparty/fish_speech/tools/server/views.py

@@ -0,0 +1,246 @@
+import io
+import os
+import time
+from http import HTTPStatus
+
+import numpy as np
+import ormsgpack
+import soundfile as sf
+import torch
+from kui.asgi import HTTPException, HttpView, JSONResponse, StreamResponse, request
+from loguru import logger
+
+from tools.schema import (
+    ServeASRRequest,
+    ServeASRResponse,
+    ServeChatRequest,
+    ServeTTSRequest,
+    ServeVQGANDecodeRequest,
+    ServeVQGANDecodeResponse,
+    ServeVQGANEncodeRequest,
+    ServeVQGANEncodeResponse,
+)
+from tools.server.agent import get_response_generator
+from tools.server.api_utils import (
+    buffer_to_async_generator,
+    get_content_type,
+    inference_async,
+)
+from tools.server.inference import inference_wrapper as inference
+from tools.server.model_manager import ModelManager
+from tools.server.model_utils import batch_asr, cached_vqgan_batch_encode, vqgan_decode
+
+MAX_NUM_SAMPLES = int(os.getenv("NUM_SAMPLES", 1))
+
+
+class HealthView(HttpView):
+    """
+    Return the health status of the server.
+    """
+
+    @classmethod
+    async def post(cls):
+        return JSONResponse({"status": "ok"})
+
+
+class VQGANEncodeView(HttpView):
+    """
+    Encode the audio into symbolic tokens.
+    """
+
+    @classmethod
+    async def post(cls):
+        # Decode the request
+        payload = await request.data()
+        req = ServeVQGANEncodeRequest(**payload)
+
+        # Get the model from the app
+        model_manager: ModelManager = request.app.state.model_manager
+        decoder_model = model_manager.decoder_model
+
+        # Encode the audio
+        start_time = time.time()
+        tokens = cached_vqgan_batch_encode(decoder_model, req.audios)
+        logger.info(
+            f"[EXEC] VQGAN encode time: {(time.time() - start_time) * 1000:.2f}ms"
+        )
+
+        # Return the response
+        return ormsgpack.packb(
+            ServeVQGANEncodeResponse(tokens=[i.tolist() for i in tokens]),
+            option=ormsgpack.OPT_SERIALIZE_PYDANTIC,
+        )
+
+
+class VQGANDecodeView(HttpView):
+    """
+    Decode the symbolic tokens into audio.
+    """
+
+    @classmethod
+    async def post(cls):
+        # Decode the request
+        payload = await request.data()
+        req = ServeVQGANDecodeRequest(**payload)
+
+        # Get the model from the app
+        model_manager: ModelManager = request.app.state.model_manager
+        decoder_model = model_manager.decoder_model
+
+        # Decode the audio
+        tokens = [torch.tensor(token, dtype=torch.int) for token in req.tokens]
+        start_time = time.time()
+        audios = vqgan_decode(decoder_model, tokens)
+        logger.info(
+            f"[EXEC] VQGAN decode time: {(time.time() - start_time) * 1000:.2f}ms"
+        )
+        audios = [audio.astype(np.float16).tobytes() for audio in audios]
+
+        # Return the response
+        return ormsgpack.packb(
+            ServeVQGANDecodeResponse(audios=audios),
+            option=ormsgpack.OPT_SERIALIZE_PYDANTIC,
+        )
+
+
+class ASRView(HttpView):
+    """
+    Perform automatic speech recognition on the audio.
+    """
+
+    @classmethod
+    async def post(cls):
+        # Decode the request
+        payload = await request.data()
+        req = ServeASRRequest(**payload)
+
+        # Get the model from the app
+        model_manager: ModelManager = request.app.state.model_manager
+        asr_model = model_manager.asr_model
+        lock = request.app.state.lock
+
+        # Perform ASR
+        start_time = time.time()
+        audios = [np.frombuffer(audio, dtype=np.float16) for audio in req.audios]
+        audios = [torch.from_numpy(audio).float() for audio in audios]
+
+        if any(audios.shape[-1] >= 30 * req.sample_rate for audios in audios):
+            raise HTTPException(status_code=400, content="Audio length is too long")
+
+        transcriptions = batch_asr(
+            asr_model, lock, audios=audios, sr=req.sample_rate, language=req.language
+        )
+        logger.info(f"[EXEC] ASR time: {(time.time() - start_time) * 1000:.2f}ms")
+
+        # Return the response
+        return ormsgpack.packb(
+            ServeASRResponse(transcriptions=transcriptions),
+            option=ormsgpack.OPT_SERIALIZE_PYDANTIC,
+        )
+
+
+class TTSView(HttpView):
+    """
+    Perform text-to-speech on the input text.
+    """
+
+    @classmethod
+    async def post(cls):
+        # Decode the request
+        payload = await request.data()
+        req = ServeTTSRequest(**payload)
+
+        # Get the model from the app
+        app_state = request.app.state
+        model_manager: ModelManager = app_state.model_manager
+        engine = model_manager.tts_inference_engine
+        sample_rate = engine.decoder_model.spec_transform.sample_rate
+
+        # Check if the text is too long
+        if app_state.max_text_length > 0 and len(req.text) > app_state.max_text_length:
+            raise HTTPException(
+                HTTPStatus.BAD_REQUEST,
+                content=f"Text is too long, max length is {app_state.max_text_length}",
+            )
+
+        # Check if streaming is enabled
+        if req.streaming and req.format != "wav":
+            raise HTTPException(
+                HTTPStatus.BAD_REQUEST,
+                content="Streaming only supports WAV format",
+            )
+
+        # Perform TTS
+        if req.streaming:
+            return StreamResponse(
+                iterable=inference_async(req, engine),
+                headers={
+                    "Content-Disposition": f"attachment; filename=audio.{req.format}",
+                },
+                content_type=get_content_type(req.format),
+            )
+        else:
+            fake_audios = next(inference(req, engine))
+            buffer = io.BytesIO()
+            sf.write(
+                buffer,
+                fake_audios,
+                sample_rate,
+                format=req.format,
+            )
+
+            return StreamResponse(
+                iterable=buffer_to_async_generator(buffer.getvalue()),
+                headers={
+                    "Content-Disposition": f"attachment; filename=audio.{req.format}",
+                },
+                content_type=get_content_type(req.format),
+            )
+
+
+class ChatView(HttpView):
+    """
+    Perform chatbot inference on the input text.
+    """
+
+    @classmethod
+    async def post(cls):
+        # Decode the request
+        payload = await request.data()
+        req = ServeChatRequest(**payload)
+
+        # Check that the number of samples requested is correct
+        if req.num_samples < 1 or req.num_samples > MAX_NUM_SAMPLES:
+            raise HTTPException(
+                HTTPStatus.BAD_REQUEST,
+                content=f"Number of samples must be between 1 and {MAX_NUM_SAMPLES}",
+            )
+
+        # Get the type of content provided
+        content_type = request.headers.get("Content-Type", "application/json")
+        json_mode = "application/json" in content_type
+
+        # Get the models from the app
+        model_manager: ModelManager = request.app.state.model_manager
+        llama_queue = model_manager.llama_queue
+        tokenizer = model_manager.tokenizer
+        config = model_manager.config
+
+        device = request.app.state.device
+
+        # Get the response generators
+        response_generator = get_response_generator(
+            llama_queue, tokenizer, config, req, device, json_mode
+        )
+
+        # Return the response in the correct format
+        if req.streaming is False:
+            result = response_generator()
+            if json_mode:
+                return JSONResponse(result.model_dump())
+            else:
+                return ormsgpack.packb(result, option=ormsgpack.OPT_SERIALIZE_PYDANTIC)
+
+        return StreamResponse(
+            iterable=response_generator(), content_type="text/event-stream"
+        )

xinference/thirdparty/fish_speech/tools/vqgan/extract_vq.py

@@ -24,6 +24,12 @@ OmegaConf.register_new_resolver("eval", eval)
 # This file is used to convert the audio files to text files using the Whisper model.
 # It's mainly used to generate the training data for the VQ model.
 
+backends = torchaudio.list_audio_backends()
+
+if "ffmpeg" in backends:
+    backend = "ffmpeg"
+else:
+    backend = "soundfile"
 
 RANK = int(os.environ.get("SLURM_PROCID", 0))
 WORLD_SIZE = int(os.environ.get("SLURM_NTASKS", 1))
@@ -81,7 +87,7 @@ def process_batch(files: list[Path], model) -> float:
     for file in files:
         try:
             wav, sr = torchaudio.load(
-                str(file), backend="sox" if sys.platform == "linux" else "soundfile"
+                str(file), backend=backend
             )  # Need to install libsox-dev
         except Exception as e:
             logger.error(f"Error reading {file}: {e}")

xinference/thirdparty/fish_speech/tools/vqgan/inference.py

@@ -24,8 +24,7 @@ def load_model(config_name, checkpoint_path, device="cuda"):
 
     model = instantiate(cfg)
     state_dict = torch.load(
-        checkpoint_path,
-        map_location=device,
+        checkpoint_path, map_location=device, mmap=True, weights_only=True
     )
     if "state_dict" in state_dict:
         state_dict = state_dict["state_dict"]
@@ -37,7 +36,7 @@ def load_model(config_name, checkpoint_path, device="cuda"):
             if "generator." in k
         }
 
-    result = model.load_state_dict(state_dict, strict=False)
+    result = model.load_state_dict(state_dict, strict=False, assign=True)
     model.eval()
     model.to(device)