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

xinference/core/utils.py

@@ -62,12 +62,16 @@ def log_async(
 
         @wraps(func)
         async def wrapped(*args, **kwargs):
-            try:
-                bound_args = sig.bind_partial(*args, **kwargs)
-                arguments = bound_args.arguments
-            except TypeError:
-                arguments = {}
-            request_id_str = arguments.get("request_id", "")
+            request_id_str = kwargs.get("request_id")
+            if not request_id_str:
+                # sometimes `request_id` not in kwargs
+                # we try to bind the arguments
+                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":
@@ -272,8 +276,8 @@ def get_nvidia_gpu_info() -> Dict:
 
 
 def assign_replica_gpu(
-    _replica_model_uid: str, replica: int, gpu_idx: Union[int, List[int]]
-) -> List[int]:
+    _replica_model_uid: str, replica: int, gpu_idx: Optional[Union[int, List[int]]]
+) -> Optional[List[int]]:
     model_uid, rep_id = parse_replica_model_uid(_replica_model_uid)
     rep_id, replica = int(rep_id), int(replica)
     if isinstance(gpu_idx, int):

xinference/core/worker.py

@@ -22,6 +22,7 @@ import signal
 import threading
 import time
 from collections import defaultdict
+from dataclasses import dataclass
 from logging import getLogger
 from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union
 
@@ -58,6 +59,11 @@ else:
     MODEL_ACTOR_AUTO_RECOVER_LIMIT = None
 
 
+@dataclass
+class ModelStatus:
+    last_error: str = ""
+
+
 class WorkerActor(xo.StatelessActor):
     def __init__(
         self,
@@ -90,6 +96,7 @@ class WorkerActor(xo.StatelessActor):
         # attributes maintained after model launched:
         self._model_uid_to_model: Dict[str, xo.ActorRefType["ModelActor"]] = {}
         self._model_uid_to_model_spec: Dict[str, ModelDescription] = {}
+        self._model_uid_to_model_status: Dict[str, ModelStatus] = {}
         self._gpu_to_model_uid: Dict[int, str] = {}
         self._gpu_to_embedding_model_uids: Dict[int, Set[str]] = defaultdict(set)
         # Dict structure: gpu_index: {(replica_model_uid, model_type)}
@@ -866,6 +873,9 @@ class WorkerActor(xo.StatelessActor):
             )
 
             try:
+                xavier_config: Optional[Dict] = kwargs.pop("xavier_config", None)
+                if xavier_config is not None:
+                    xavier_config["rank_address"] = subpool_address
                 model, model_description = await asyncio.to_thread(
                     create_model_instance,
                     subpool_address,
@@ -893,6 +903,7 @@ class WorkerActor(xo.StatelessActor):
                     model=model,
                     model_description=model_description,
                     request_limits=request_limits,
+                    xavier_config=xavier_config,
                 )
                 await model_ref.load()
             except:
@@ -902,6 +913,7 @@ class WorkerActor(xo.StatelessActor):
                 raise
             self._model_uid_to_model[model_uid] = model_ref
             self._model_uid_to_model_spec[model_uid] = model_description
+            self._model_uid_to_model_status[model_uid] = ModelStatus()
             self._model_uid_to_addr[model_uid] = subpool_address
             self._model_uid_to_recover_count.setdefault(
                 model_uid, MODEL_ACTOR_AUTO_RECOVER_LIMIT
@@ -921,6 +933,7 @@ class WorkerActor(xo.StatelessActor):
             origin_uid,
             {"model_ability": abilities, "status": LaunchStatus.READY.name},
         )
+        return subpool_address
 
     @log_async(logger=logger, level=logging.INFO)
     async def terminate_model(self, model_uid: str, is_model_die=False):
@@ -976,6 +989,7 @@ class WorkerActor(xo.StatelessActor):
                 status = LaunchStatus.ERROR.name
             else:
                 status = LaunchStatus.TERMINATED.name
+                self._model_uid_to_model_status.pop(model_uid, None)
 
             if self._status_guard_ref is None:
                 _ = await self.get_supervisor_ref()
@@ -1010,6 +1024,9 @@ class WorkerActor(xo.StatelessActor):
 
     @log_sync(logger=logger)
     def get_model(self, model_uid: str) -> xo.ActorRefType["ModelActor"]:
+        model_status = self._model_uid_to_model_status.get(model_uid)
+        if model_status and model_status.last_error:
+            raise Exception(model_status.last_error)
         model_ref = self._model_uid_to_model.get(model_uid, None)
         if model_ref is None:
             raise ValueError(f"Model not found, uid: {model_uid}")
@@ -1138,6 +1155,21 @@ class WorkerActor(xo.StatelessActor):
         }
         return ret
 
+    def update_model_status(self, model_uid: str, **kwargs):
+        model_status = self._model_uid_to_model_status.get(model_uid)
+        if model_status is not None:
+            for k, v in kwargs.items():
+                setattr(model_status, k, v)
+
+    def get_model_status(self, model_uid: str):
+        return self._model_uid_to_model_status.get(model_uid)
+
     @staticmethod
     def record_metrics(name, op, kwargs):
         record_metrics(name, op, kwargs)
+
+    async def start_transfer_for_vllm(
+        self, rep_model_uid: str, rank_addresses: List[str]
+    ):
+        model_ref = self._model_uid_to_model[rep_model_uid]
+        await model_ref.start_transfer_for_vllm(rank_addresses)

xinference/model/audio/core.py

@@ -22,6 +22,7 @@ from ..utils import valid_model_revision
 from .chattts import ChatTTSModel
 from .cosyvoice import CosyVoiceModel
 from .f5tts import F5TTSModel
+from .f5tts_mlx import F5TTSMLXModel
 from .fish_speech import FishSpeechModel
 from .funasr import FunASRModel
 from .whisper import WhisperModel
@@ -171,6 +172,7 @@ def create_audio_model_instance(
         CosyVoiceModel,
         FishSpeechModel,
         F5TTSModel,
+        F5TTSMLXModel,
     ],
     AudioModelDescription,
 ]:
@@ -185,6 +187,7 @@ def create_audio_model_instance(
         CosyVoiceModel,
         FishSpeechModel,
         F5TTSModel,
+        F5TTSMLXModel,
     ]
     if model_spec.model_family == "whisper":
         if not model_spec.engine:
@@ -201,6 +204,8 @@ def create_audio_model_instance(
         model = FishSpeechModel(model_uid, model_path, model_spec, **kwargs)
     elif model_spec.model_family == "F5-TTS":
         model = F5TTSModel(model_uid, model_path, model_spec, **kwargs)
+    elif model_spec.model_family == "F5-TTS-MLX":
+        model = F5TTSMLXModel(model_uid, model_path, model_spec, **kwargs)
     else:
         raise Exception(f"Unsupported audio model family: {model_spec.model_family}")
     model_description = AudioModelDescription(

xinference/model/audio/cosyvoice.py

@@ -39,6 +39,7 @@ class CosyVoiceModel:
         self._device = device
         self._model = None
         self._kwargs = kwargs
+        self._is_cosyvoice2 = False
 
     @property
     def model_ability(self):
@@ -51,7 +52,14 @@ class CosyVoiceModel:
         # The yaml config loaded from model has hard-coded the import paths. please refer to: load_hyperpyyaml
         sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../thirdparty"))
 
-        from cosyvoice.cli.cosyvoice import CosyVoice
+        if "CosyVoice2" in self._model_spec.model_name:
+            from cosyvoice.cli.cosyvoice import CosyVoice2 as CosyVoice
+
+            self._is_cosyvoice2 = True
+        else:
+            from cosyvoice.cli.cosyvoice import CosyVoice
+
+            self._is_cosyvoice2 = False
 
         self._model = CosyVoice(
             self._model_path, load_jit=self._kwargs.get("load_jit", False)
@@ -78,12 +86,22 @@ class CosyVoiceModel:
                 output = self._model.inference_zero_shot(
                     input, prompt_text, prompt_speech_16k, stream=stream
                 )
+            elif instruct_text:
+                assert self._is_cosyvoice2
+                logger.info("CosyVoice inference_instruct")
+                output = self._model.inference_instruct2(
+                    input,
+                    instruct_text=instruct_text,
+                    prompt_speech_16k=prompt_speech_16k,
+                    stream=stream,
+                )
             else:
                 logger.info("CosyVoice inference_cross_lingual")
                 output = self._model.inference_cross_lingual(
                     input, prompt_speech_16k, stream=stream
                 )
         else:
+            assert not self._is_cosyvoice2
             available_speakers = self._model.list_avaliable_spks()
             if not voice:
                 voice = available_speakers[0]
@@ -106,7 +124,9 @@ class CosyVoiceModel:
         def _generator_stream():
             with BytesIO() as out:
                 writer = torchaudio.io.StreamWriter(out, format=response_format)
-                writer.add_audio_stream(sample_rate=22050, num_channels=1)
+                writer.add_audio_stream(
+                    sample_rate=self._model.sample_rate, num_channels=1
+                )
                 i = 0
                 last_pos = 0
                 with writer.open():
@@ -125,7 +145,7 @@ class CosyVoiceModel:
             chunks = [o["tts_speech"] for o in output]
             t = torch.cat(chunks, dim=1)
             with BytesIO() as out:
-                torchaudio.save(out, t, 22050, format=response_format)
+                torchaudio.save(out, t, self._model.sample_rate, format=response_format)
                 return out.getvalue()
 
         return _generator_stream() if stream else _generator_block()
@@ -163,6 +183,8 @@ class CosyVoiceModel:
             assert (
                 prompt_text is None
             ), "CosyVoice Instruct model does not support prompt_text"
+        elif self._is_cosyvoice2:
+            assert prompt_speech is not None, "CosyVoice2 requires prompt_speech"
         else:
             # inference_zero_shot
             # inference_cross_lingual

xinference/model/audio/f5tts.py

@@ -11,12 +11,12 @@
 # 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 io
 import logging
 import os
 import re
 from io import BytesIO
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Union
 
 if TYPE_CHECKING:
     from .core import AudioModelFamilyV1
@@ -106,9 +106,9 @@ class F5TTSModel:
             ) = preprocess_ref_audio_text(
                 voices[voice]["ref_audio"], voices[voice]["ref_text"]
             )
-            print("Voice:", voice)
-            print("Ref_audio:", voices[voice]["ref_audio"])
-            print("Ref_text:", voices[voice]["ref_text"])
+            logger.info("Voice:", voice)
+            logger.info("Ref_audio:", voices[voice]["ref_audio"])
+            logger.info("Ref_text:", voices[voice]["ref_text"])
 
         final_sample_rate = None
         generated_audio_segments = []
@@ -122,16 +122,16 @@ class F5TTSModel:
             if match:
                 voice = match[1]
             else:
-                print("No voice tag found, using main.")
+                logger.info("No voice tag found, using main.")
                 voice = "main"
             if voice not in voices:
-                print(f"Voice {voice} not found, using main.")
+                logger.info(f"Voice {voice} not found, using main.")
                 voice = "main"
             text = re.sub(reg2, "", text)
             gen_text = text.strip()
             ref_audio = voices[voice]["ref_audio"]
             ref_text = voices[voice]["ref_text"]
-            print(f"Voice: {voice}")
+            logger.info(f"Voice: {voice}")
             audio, final_sample_rate, spectragram = infer_process(
                 ref_audio,
                 ref_text,
@@ -167,18 +167,23 @@ class F5TTSModel:
         prompt_speech: Optional[bytes] = kwargs.pop("prompt_speech", None)
         prompt_text: Optional[str] = kwargs.pop("prompt_text", None)
 
+        ref_audio: Union[str, io.BytesIO]
         if prompt_speech is None:
             base = os.path.dirname(f5_tts.__file__)
             config = os.path.join(base, "infer/examples/basic/basic.toml")
             with open(config, "rb") as f:
                 config_dict = tomli.load(f)
-            prompt_speech = os.path.join(base, config_dict["ref_audio"])
+            ref_audio = os.path.join(base, config_dict["ref_audio"])
             prompt_text = config_dict["ref_text"]
+        else:
+            ref_audio = io.BytesIO(prompt_speech)
+            if prompt_text is None:
+                raise ValueError("`prompt_text` cannot be empty")
 
         assert self._model is not None
         vocoder_name = self._kwargs.get("vocoder_name", "vocos")
         sample_rate, wav = self._infer(
-            ref_audio=prompt_speech,
+            ref_audio=ref_audio,
             ref_text=prompt_text,
             text_gen=input,
             model_obj=self._model,

xinference/model/audio/f5tts_mlx.py

@@ -0,0 +1,260 @@
+# 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 datetime
+import io
+import logging
+import os
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, Optional, Union
+
+import numpy as np
+from tqdm import tqdm
+
+if TYPE_CHECKING:
+    from .core import AudioModelFamilyV1
+
+logger = logging.getLogger(__name__)
+
+
+class F5TTSMLXModel:
+    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._model = None
+
+    @property
+    def model_ability(self):
+        return self._model_spec.model_ability
+
+    def load(self):
+        try:
+            import mlx.core as mx
+            from f5_tts_mlx.cfm import F5TTS
+            from f5_tts_mlx.dit import DiT
+            from f5_tts_mlx.duration import DurationPredictor, DurationTransformer
+            from vocos_mlx import Vocos
+        except ImportError:
+            error_message = "Failed to import module 'f5_tts_mlx'"
+            installation_guide = [
+                "Please make sure 'f5_tts_mlx' is installed.\n",
+            ]
+
+            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
+
+        path = Path(self._model_path)
+        # vocab
+
+        vocab_path = path / "vocab.txt"
+        vocab = {v: i for i, v in enumerate(Path(vocab_path).read_text().split("\n"))}
+        if len(vocab) == 0:
+            raise ValueError(f"Could not load vocab from {vocab_path}")
+
+        # duration predictor
+
+        duration_model_path = path / "duration_v2.safetensors"
+        duration_predictor = None
+
+        if duration_model_path.exists():
+            duration_predictor = DurationPredictor(
+                transformer=DurationTransformer(
+                    dim=512,
+                    depth=8,
+                    heads=8,
+                    text_dim=512,
+                    ff_mult=2,
+                    conv_layers=2,
+                    text_num_embeds=len(vocab) - 1,
+                ),
+                vocab_char_map=vocab,
+            )
+            weights = mx.load(duration_model_path.as_posix(), format="safetensors")
+            duration_predictor.load_weights(list(weights.items()))
+
+        # vocoder
+
+        vocos = Vocos.from_pretrained("lucasnewman/vocos-mel-24khz")
+
+        # model
+
+        model_path = path / "model.safetensors"
+
+        f5tts = F5TTS(
+            transformer=DiT(
+                dim=1024,
+                depth=22,
+                heads=16,
+                ff_mult=2,
+                text_dim=512,
+                conv_layers=4,
+                text_num_embeds=len(vocab) - 1,
+            ),
+            vocab_char_map=vocab,
+            vocoder=vocos.decode,
+            duration_predictor=duration_predictor,
+        )
+
+        weights = mx.load(model_path.as_posix(), format="safetensors")
+        f5tts.load_weights(list(weights.items()))
+        mx.eval(f5tts.parameters())
+
+        self._model = f5tts
+
+    def speech(
+        self,
+        input: str,
+        voice: str,
+        response_format: str = "mp3",
+        speed: float = 1.0,
+        stream: bool = False,
+        **kwargs,
+    ):
+        import mlx.core as mx
+        import soundfile as sf
+        import tomli
+        from f5_tts_mlx.generate import (
+            FRAMES_PER_SEC,
+            SAMPLE_RATE,
+            TARGET_RMS,
+            convert_char_to_pinyin,
+            split_sentences,
+        )
+
+        from .utils import ensure_sample_rate
+
+        if stream:
+            raise Exception("F5-TTS does not support stream generation.")
+
+        prompt_speech: Optional[bytes] = kwargs.pop("prompt_speech", None)
+        prompt_text: Optional[str] = kwargs.pop("prompt_text", None)
+        duration: Optional[float] = kwargs.pop("duration", None)
+        steps: Optional[int] = kwargs.pop("steps", 8)
+        cfg_strength: Optional[float] = kwargs.pop("cfg_strength", 2.0)
+        method: Literal["euler", "midpoint"] = kwargs.pop("method", "rk4")
+        sway_sampling_coef: float = kwargs.pop("sway_sampling_coef", -1.0)
+        seed: Optional[int] = kwargs.pop("seed", None)
+
+        prompt_speech_path: Union[str, io.BytesIO]
+        if prompt_speech is None:
+            base = os.path.join(os.path.dirname(__file__), "../../thirdparty/f5_tts")
+            config = os.path.join(base, "infer/examples/basic/basic.toml")
+            with open(config, "rb") as f:
+                config_dict = tomli.load(f)
+            prompt_speech_path = os.path.join(base, config_dict["ref_audio"])
+            prompt_text = config_dict["ref_text"]
+        else:
+            prompt_speech_path = io.BytesIO(prompt_speech)
+
+            if prompt_text is None:
+                raise ValueError("`prompt_text` cannot be empty")
+
+        audio, sr = sf.read(prompt_speech_path)
+        audio = ensure_sample_rate(audio, sr, SAMPLE_RATE)
+
+        audio = mx.array(audio)
+        ref_audio_duration = audio.shape[0] / SAMPLE_RATE
+        logger.debug(
+            f"Got reference audio with duration: {ref_audio_duration:.2f} seconds"
+        )
+
+        rms = mx.sqrt(mx.mean(mx.square(audio)))
+        if rms < TARGET_RMS:
+            audio = audio * TARGET_RMS / rms
+
+        sentences = split_sentences(input)
+        is_single_generation = len(sentences) <= 1 or duration is not None
+
+        if is_single_generation:
+            generation_text = convert_char_to_pinyin([prompt_text + " " + input])  # type: ignore
+
+            if duration is not None:
+                duration = int(duration * FRAMES_PER_SEC)
+
+            start_date = datetime.datetime.now()
+
+            wave, _ = self._model.sample(  # type: ignore
+                mx.expand_dims(audio, axis=0),
+                text=generation_text,
+                duration=duration,
+                steps=steps,
+                method=method,
+                speed=speed,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+                seed=seed,
+            )
+
+            wave = wave[audio.shape[0] :]
+            mx.eval(wave)
+
+            generated_duration = wave.shape[0] / SAMPLE_RATE
+            print(
+                f"Generated {generated_duration:.2f}s of audio in {datetime.datetime.now() - start_date}."
+            )
+
+        else:
+            start_date = datetime.datetime.now()
+
+            output = []
+
+            for sentence_text in tqdm(split_sentences(input)):
+                text = convert_char_to_pinyin([prompt_text + " " + sentence_text])  # type: ignore
+
+                if duration is not None:
+                    duration = int(duration * FRAMES_PER_SEC)
+
+                wave, _ = self._model.sample(  # type: ignore
+                    mx.expand_dims(audio, axis=0),
+                    text=text,
+                    duration=duration,
+                    steps=steps,
+                    method=method,
+                    speed=speed,
+                    cfg_strength=cfg_strength,
+                    sway_sampling_coef=sway_sampling_coef,
+                    seed=seed,
+                )
+
+                # trim the reference audio
+                wave = wave[audio.shape[0] :]
+                mx.eval(wave)
+
+                output.append(wave)
+
+            wave = mx.concatenate(output, axis=0)
+
+            generated_duration = wave.shape[0] / SAMPLE_RATE
+            logger.debug(
+                f"Generated {generated_duration:.2f}s of audio in {datetime.datetime.now() - start_date}."
+            )
+
+        # Save the generated audio
+        with BytesIO() as out:
+            with sf.SoundFile(
+                out, "w", SAMPLE_RATE, 1, format=response_format.upper()
+            ) as f:
+                f.write(np.array(wave))
+                return out.getvalue()