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

xinference/core/utils.py

@@ -11,11 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import asyncio
 import logging
 import os
 import random
 import string
 import uuid
+import weakref
 from enum import Enum
 from typing import Dict, Generator, List, Optional, Tuple, Union
 
@@ -23,7 +25,10 @@ import orjson
 from pynvml import nvmlDeviceGetCount, nvmlInit, nvmlShutdown
 
 from .._compat import BaseModel
-from ..constants import XINFERENCE_LOG_ARG_MAX_LENGTH
+from ..constants import (
+    XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
+    XINFERENCE_LOG_ARG_MAX_LENGTH,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -49,13 +54,20 @@ def log_async(
 ):
     import time
     from functools import wraps
+    from inspect import signature
 
     def decorator(func):
         func_name = func.__name__
+        sig = signature(func)
 
         @wraps(func)
         async def wrapped(*args, **kwargs):
-            request_id_str = kwargs.get("request_id", "")
+            try:
+                bound_args = sig.bind_partial(*args, **kwargs)
+                arguments = bound_args.arguments
+            except TypeError:
+                arguments = {}
+            request_id_str = arguments.get("request_id", "")
             if not request_id_str:
                 request_id_str = uuid.uuid1()
                 if func_name == "text_to_image":
@@ -269,3 +281,56 @@ def assign_replica_gpu(
     if isinstance(gpu_idx, list) and gpu_idx:
         return gpu_idx[rep_id::replica]
     return gpu_idx
+
+
+class CancelMixin:
+    _CANCEL_TASK_NAME = "abort_block"
+
+    def __init__(self):
+        self._running_tasks: weakref.WeakValueDictionary[
+            str, asyncio.Task
+        ] = weakref.WeakValueDictionary()
+
+    def _add_running_task(self, request_id: Optional[str]):
+        """Add current asyncio task to the running task.
+        :param request_id: The corresponding request id.
+        """
+        if request_id is None:
+            return
+        running_task = self._running_tasks.get(request_id)
+        if running_task is not None:
+            if running_task.get_name() == self._CANCEL_TASK_NAME:
+                raise Exception(f"The request has been aborted: {request_id}")
+            raise Exception(f"Duplicate request id: {request_id}")
+        current_task = asyncio.current_task()
+        assert current_task is not None
+        self._running_tasks[request_id] = current_task
+
+    def _cancel_running_task(
+        self,
+        request_id: Optional[str],
+        block_duration: int = XINFERENCE_DEFAULT_CANCEL_BLOCK_DURATION,
+    ):
+        """Cancel the running asyncio task.
+        :param request_id: The request id to cancel.
+        :param block_duration: The duration seconds to ensure the request can't be executed.
+        """
+        if request_id is None:
+            return
+        running_task = self._running_tasks.pop(request_id, None)
+        if running_task is not None:
+            running_task.cancel()
+
+        async def block_task():
+            """This task is for blocking the request for a duration."""
+            try:
+                await asyncio.sleep(block_duration)
+                logger.info("Abort block end for request: %s", request_id)
+            except asyncio.CancelledError:
+                logger.info("Abort block is cancelled for request: %s", request_id)
+
+        if block_duration > 0:
+            logger.info("Abort block start for request: %s", request_id)
+            self._running_tasks[request_id] = asyncio.create_task(
+                block_task(), name=self._CANCEL_TASK_NAME
+            )

xinference/model/audio/__init__.py

@@ -15,6 +15,8 @@
 import codecs
 import json
 import os
+import platform
+import sys
 import warnings
 from typing import Any, Dict
 
@@ -55,6 +57,14 @@ def register_custom_model():
                 warnings.warn(f"{user_defined_audio_dir}/{f} has error, {e}")
 
 
+def _need_filter(spec: dict):
+    if (sys.platform != "darwin" or platform.processor() != "arm") and spec.get(
+        "engine", ""
+    ).upper() == "MLX":
+        return True
+    return False
+
+
 def _install():
     _model_spec_json = os.path.join(os.path.dirname(__file__), "model_spec.json")
     _model_spec_modelscope_json = os.path.join(
@@ -64,6 +74,7 @@ def _install():
         dict(
             (spec["model_name"], AudioModelFamilyV1(**spec))
             for spec in json.load(codecs.open(_model_spec_json, "r", encoding="utf-8"))
+            if not _need_filter(spec)
         )
     )
     for model_name, model_spec in BUILTIN_AUDIO_MODELS.items():
@@ -75,6 +86,7 @@ def _install():
             for spec in json.load(
                 codecs.open(_model_spec_modelscope_json, "r", encoding="utf-8")
             )
+            if not _need_filter(spec)
         )
     )
     for model_name, model_spec in MODELSCOPE_AUDIO_MODELS.items():

xinference/model/audio/core.py

@@ -24,6 +24,7 @@ from .cosyvoice import CosyVoiceModel
 from .fish_speech import FishSpeechModel
 from .funasr import FunASRModel
 from .whisper import WhisperModel
+from .whisper_mlx import WhisperMLXModel
 
 logger = logging.getLogger(__name__)
 
@@ -43,11 +44,12 @@ class AudioModelFamilyV1(CacheableModelSpec):
     model_family: str
     model_name: str
     model_id: str
-    model_revision: str
+    model_revision: Optional[str]
     multilingual: bool
     model_ability: Optional[str]
     default_model_config: Optional[Dict[str, Any]]
     default_transcription_config: Optional[Dict[str, Any]]
+    engine: Optional[str]
 
 
 class AudioModelDescription(ModelDescription):
@@ -160,17 +162,32 @@ def create_audio_model_instance(
     model_path: Optional[str] = None,
     **kwargs,
 ) -> Tuple[
-    Union[WhisperModel, FunASRModel, ChatTTSModel, CosyVoiceModel, FishSpeechModel],
+    Union[
+        WhisperModel,
+        WhisperMLXModel,
+        FunASRModel,
+        ChatTTSModel,
+        CosyVoiceModel,
+        FishSpeechModel,
+    ],
     AudioModelDescription,
 ]:
     model_spec = match_audio(model_name, download_hub)
     if model_path is None:
         model_path = cache(model_spec)
     model: Union[
-        WhisperModel, FunASRModel, ChatTTSModel, CosyVoiceModel, FishSpeechModel
+        WhisperModel,
+        WhisperMLXModel,
+        FunASRModel,
+        ChatTTSModel,
+        CosyVoiceModel,
+        FishSpeechModel,
     ]
     if model_spec.model_family == "whisper":
-        model = WhisperModel(model_uid, model_path, model_spec, **kwargs)
+        if not model_spec.engine:
+            model = WhisperModel(model_uid, model_path, model_spec, **kwargs)
+        else:
+            model = WhisperMLXModel(model_uid, model_path, model_spec, **kwargs)
     elif model_spec.model_family == "funasr":
         model = FunASRModel(model_uid, model_path, model_spec, **kwargs)
     elif model_spec.model_family == "ChatTTS":

xinference/model/audio/fish_speech.py

@@ -81,12 +81,14 @@ class FishSpeechModel:
             if not is_device_available(self._device):
                 raise ValueError(f"Device {self._device} is not available!")
 
-        logger.info("Loading Llama model...")
+        enable_compile = self._kwargs.get("compile", False)
+        precision = self._kwargs.get("precision", torch.bfloat16)
+        logger.info("Loading Llama model, compile=%s...", enable_compile)
         self._llama_queue = launch_thread_safe_queue(
             checkpoint_path=self._model_path,
             device=self._device,
-            precision=torch.bfloat16,
-            compile=False,
+            precision=precision,
+            compile=enable_compile,
         )
         logger.info("Llama model loaded, loading VQ-GAN model...")
 
@@ -112,9 +114,10 @@ class FishSpeechModel:
         top_p,
         repetition_penalty,
         temperature,
+        seed="0",
         streaming=False,
     ):
-        from fish_speech.utils import autocast_exclude_mps
+        from fish_speech.utils import autocast_exclude_mps, set_seed
         from tools.api import decode_vq_tokens, encode_reference
         from tools.llama.generate import (
             GenerateRequest,
@@ -122,6 +125,11 @@ class FishSpeechModel:
             WrappedGenerateResponse,
         )
 
+        seed = int(seed)
+        if seed != 0:
+            set_seed(seed)
+            logger.warning(f"set seed: {seed}")
+
         # Parse reference audio aka prompt
         prompt_tokens = encode_reference(
             decoder_model=self._model,
@@ -137,7 +145,7 @@ class FishSpeechModel:
             top_p=top_p,
             repetition_penalty=repetition_penalty,
             temperature=temperature,
-            compile=False,
+            compile=self._kwargs.get("compile", False),
             iterative_prompt=chunk_length > 0,
             chunk_length=chunk_length,
             max_length=2048,
@@ -153,22 +161,20 @@ class FishSpeechModel:
             )
         )
 
-        if streaming:
-            yield wav_chunk_header(), None, None
-
         segments = []
 
         while True:
-            result: WrappedGenerateResponse = response_queue.get()  # type: ignore
+            result: WrappedGenerateResponse = response_queue.get()
             if result.status == "error":
-                raise Exception(str(result.response))
+                raise result.response
 
-            result: GenerateResponse = result.response  # type: ignore
+            result: GenerateResponse = result.response
             if result.action == "next":
                 break
 
             with autocast_exclude_mps(
-                device_type=self._model.device.type, dtype=torch.bfloat16
+                device_type=self._model.device.type,
+                dtype=self._kwargs.get("precision", torch.bfloat16),
             ):
                 fake_audios = decode_vq_tokens(
                     decoder_model=self._model,
@@ -179,7 +185,7 @@ class FishSpeechModel:
             segments.append(fake_audios)
 
             if streaming:
-                yield (fake_audios * 32768).astype(np.int16).tobytes(), None, None
+                yield fake_audios, None, None
 
         if len(segments) == 0:
             raise Exception("No audio generated, please check the input text.")
@@ -204,29 +210,58 @@ class FishSpeechModel:
         logger.warning("Fish speech does not support setting voice: %s.", voice)
         if speed != 1.0:
             logger.warning("Fish speech does not support setting speed: %s.", speed)
-        if stream is True:
-            logger.warning("stream mode is not implemented.")
         import torchaudio
 
-        result = list(
-            self._inference(
-                text=input,
-                enable_reference_audio=False,
-                reference_audio=None,
-                reference_text=kwargs.get("reference_text", ""),
-                max_new_tokens=kwargs.get("max_new_tokens", 1024),
-                chunk_length=kwargs.get("chunk_length", 200),
-                top_p=kwargs.get("top_p", 0.7),
-                repetition_penalty=kwargs.get("repetition_penalty", 1.2),
-                temperature=kwargs.get("temperature", 0.7),
-            )
+        prompt_speech = kwargs.get("prompt_speech")
+        result = self._inference(
+            text=input,
+            enable_reference_audio=kwargs.get(
+                "enable_reference_audio", prompt_speech is not None
+            ),
+            reference_audio=prompt_speech,
+            reference_text=kwargs.get("reference_text", ""),
+            max_new_tokens=kwargs.get("max_new_tokens", 1024),
+            chunk_length=kwargs.get("chunk_length", 200),
+            top_p=kwargs.get("top_p", 0.7),
+            repetition_penalty=kwargs.get("repetition_penalty", 1.2),
+            temperature=kwargs.get("temperature", 0.7),
+            streaming=stream,
         )
-        sample_rate, audio = result[0][1]
-        audio = np.array([audio])
 
-        # Save the generated audio
-        with BytesIO() as out:
-            torchaudio.save(
-                out, torch.from_numpy(audio), sample_rate, format=response_format
-            )
-            return out.getvalue()
+        if stream:
+
+            def _stream_generator():
+                with BytesIO() as out:
+                    writer = torchaudio.io.StreamWriter(out, format=response_format)
+                    writer.add_audio_stream(
+                        sample_rate=self._model.spec_transform.sample_rate,
+                        num_channels=1,
+                    )
+                    i = 0
+                    last_pos = 0
+                    with writer.open():
+                        for chunk in result:
+                            chunk = chunk[0]
+                            if chunk is not None:
+                                chunk = chunk.reshape((chunk.shape[0], 1))
+                                trans_chunk = torch.from_numpy(chunk)
+                                writer.write_audio_chunk(i, trans_chunk)
+                                new_last_pos = out.tell()
+                                if new_last_pos != last_pos:
+                                    out.seek(last_pos)
+                                    encoded_bytes = out.read()
+                                    yield encoded_bytes
+                                    last_pos = new_last_pos
+
+            return _stream_generator()
+        else:
+            result = list(result)
+            sample_rate, audio = result[0][1]
+            audio = np.array([audio])
+
+            # Save the generated audio
+            with BytesIO() as out:
+                torchaudio.save(
+                    out, torch.from_numpy(audio), sample_rate, format=response_format
+                )
+                return out.getvalue()

xinference/model/audio/model_spec.json

@@ -103,6 +103,86 @@
     "model_ability": "audio-to-text",
     "multilingual": false
   },
+  {
+    "model_name": "whisper-tiny-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-tiny",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-tiny.en-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-tiny.en-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": false,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-base-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-base-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-base.en-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-base.en-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": false,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-small-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-small-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-small.en-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-small.en-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": false,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-medium-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-medium-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-medium.en-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-medium.en-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": false,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-large-v3-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-large-v3-mlx",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
+  {
+    "model_name": "whisper-large-v3-turbo-mlx",
+    "model_family": "whisper",
+    "model_id": "mlx-community/whisper-large-v3-turbo",
+    "model_ability": "audio-to-text",
+    "multilingual": true,
+    "engine": "mlx"
+  },
   {
     "model_name": "SenseVoiceSmall",
     "model_family": "funasr",
@@ -159,7 +239,7 @@
     "model_name": "FishSpeech-1.4",
     "model_family": "FishAudio",
     "model_id": "fishaudio/fish-speech-1.4",
-    "model_revision": "3c49651b8e583b6b13f55e375432e0d57e1aa84d",
+    "model_revision": "069c573759936b35191d3380deb89183c0656f59",
     "model_ability": "text-to-audio",
     "multilingual": true
   }

xinference/model/audio/whisper_mlx.py

@@ -0,0 +1,208 @@
+# Copyright 2022-2023 XProbe Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import functools
+import itertools
+import logging
+import tempfile
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from .core import AudioModelFamilyV1
+
+logger = logging.getLogger(__name__)
+
+
+class WhisperMLXModel:
+    def __init__(
+        self,
+        model_uid: str,
+        model_path: str,
+        model_spec: "AudioModelFamilyV1",
+        device: Optional[str] = None,
+        **kwargs,
+    ):
+        self._model_uid = model_uid
+        self._model_path = model_path
+        self._model_spec = model_spec
+        self._device = device
+        self._model = None
+        self._kwargs = kwargs
+        self._use_lighting = False
+
+    @property
+    def model_ability(self):
+        return self._model_spec.model_ability
+
+    def load(self):
+        use_lightning = self._kwargs.get("use_lightning", "auto")
+        if use_lightning not in ("auto", True, False, None):
+            raise ValueError("use_lightning can only be True, False, None or auto")
+
+        if use_lightning == "auto" or use_lightning is True:
+            try:
+                import mlx.core as mx
+                from lightning_whisper_mlx.transcribe import ModelHolder
+            except ImportError:
+                if use_lightning == "auto":
+                    use_lightning = False
+                else:
+                    error_message = "Failed to import module 'lightning_whisper_mlx'"
+                    installation_guide = [
+                        "Please make sure 'lightning_whisper_mlx' is installed.\n",
+                    ]
+
+                    raise ImportError(
+                        f"{error_message}\n\n{''.join(installation_guide)}"
+                    )
+            else:
+                use_lightning = True
+        if not use_lightning:
+            try:
+                import mlx.core as mx  # noqa: F811
+                from mlx_whisper.transcribe import ModelHolder  # noqa: F811
+            except ImportError:
+                error_message = "Failed to import module 'mlx_whisper'"
+                installation_guide = [
+                    "Please make sure 'mlx_whisper' is installed.\n",
+                ]
+
+                raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
+            else:
+                use_lightning = False
+
+        logger.info(
+            "Loading MLX whisper from %s, use lightning: %s",
+            self._model_path,
+            use_lightning,
+        )
+        self._use_lighting = use_lightning
+        self._model = ModelHolder.get_model(self._model_path, mx.float16)
+
+    def transcriptions(
+        self,
+        audio: bytes,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        temperature: float = 0,
+        timestamp_granularities: Optional[List[str]] = None,
+    ):
+        return self._call(
+            audio,
+            language=language,
+            prompt=prompt,
+            response_format=response_format,
+            temperature=temperature,
+            timestamp_granularities=timestamp_granularities,
+            task="transcribe",
+        )
+
+    def translations(
+        self,
+        audio: bytes,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        temperature: float = 0,
+        timestamp_granularities: Optional[List[str]] = None,
+    ):
+        if not self._model_spec.multilingual:
+            raise RuntimeError(
+                f"Model {self._model_spec.model_name} is not suitable for translations."
+            )
+        return self._call(
+            audio,
+            language=language,
+            prompt=prompt,
+            response_format=response_format,
+            temperature=temperature,
+            timestamp_granularities=timestamp_granularities,
+            task="translate",
+        )
+
+    def _call(
+        self,
+        audio: bytes,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        temperature: float = 0,
+        timestamp_granularities: Optional[List[str]] = None,
+        task: str = "transcribe",
+    ):
+        if self._use_lighting:
+            from lightning_whisper_mlx.transcribe import transcribe_audio
+
+            transcribe = functools.partial(
+                transcribe_audio, batch_size=self._kwargs.get("batch_size", 12)
+            )
+        else:
+            from mlx_whisper import transcribe  # type: ignore
+
+        with tempfile.NamedTemporaryFile(delete=True) as f:
+            f.write(audio)
+
+            kwargs = {"task": task}
+            if response_format == "verbose_json":
+                if timestamp_granularities == ["word"]:
+                    kwargs["word_timestamps"] = True  # type: ignore
+
+            result = transcribe(
+                f.name,
+                path_or_hf_repo=self._model_path,
+                language=language,
+                temperature=temperature,
+                initial_prompt=prompt,
+                **kwargs,
+            )
+            text = result["text"]
+            segments = result["segments"]
+            language = result["language"]
+
+            if response_format == "json":
+                return {"text": text}
+            elif response_format == "verbose_json":
+                if not timestamp_granularities or timestamp_granularities == [
+                    "segment"
+                ]:
+                    return {
+                        "task": task,
+                        "language": language,
+                        "duration": segments[-1]["end"] if segments else 0,
+                        "text": text,
+                        "segments": segments,
+                    }
+                else:
+                    assert timestamp_granularities == ["word"]
+
+                    def _extract_word(word: dict) -> dict:
+                        return {
+                            "start": word["start"].item(),
+                            "end": word["end"].item(),
+                            "word": word["word"],
+                        }
+
+                    words = [
+                        _extract_word(w)
+                        for w in itertools.chain(*[s["words"] for s in segments])
+                    ]
+                    return {
+                        "task": task,
+                        "language": language,
+                        "duration": words[-1]["end"] if words else 0,
+                        "text": text,
+                        "words": words,
+                    }
+            else:
+                raise ValueError(f"Unsupported response format: {response_format}")