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

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):
@@ -48,14 +49,32 @@ class CosyVoiceModel:
         import os
         import sys
 
+        import torch
+
         # 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"))
+        thirdparty_dir = os.path.join(os.path.dirname(__file__), "../../thirdparty")
+        sys.path.insert(0, thirdparty_dir)
+
+        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
 
-        from cosyvoice.cli.cosyvoice import CosyVoice
+            self._is_cosyvoice2 = False
 
-        self._model = CosyVoice(
-            self._model_path, load_jit=self._kwargs.get("load_jit", False)
+        # Unify this configuration name as 'compile' to be compatible with the name 'load_jit'.
+        load_jit = self._kwargs.get("load_jit", False) or self._kwargs.get(
+            "compile", False
         )
+        logger.info("Loading CosyVoice model, compile=%s...", load_jit)
+        self._model = CosyVoice(self._model_path, load_jit=load_jit)
+        if self._is_cosyvoice2:
+            spk2info_file = os.path.join(thirdparty_dir, "cosyvoice/bin/spk2info.pt")
+            self._model.frontend.spk2info = torch.load(
+                spk2info_file, map_location=self._device
+            )
 
     def _speech_handle(
         self,
@@ -78,6 +97,15 @@ 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(
@@ -87,6 +115,7 @@ class CosyVoiceModel:
             available_speakers = self._model.list_avaliable_spks()
             if not voice:
                 voice = available_speakers[0]
+                logger.info("Auto select speaker: %s", voice)
             else:
                 assert (
                     voice in available_speakers
@@ -106,7 +135,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 +156,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 +194,8 @@ class CosyVoiceModel:
             assert (
                 prompt_text is None
             ), "CosyVoice Instruct model does not support prompt_text"
+        elif self._is_cosyvoice2:
+            pass
         else:
             # inference_zero_shot
             # inference_cross_lingual

xinference/model/audio/f5tts.py

@@ -0,0 +1,200 @@
+# 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 io
+import logging
+import os
+import re
+from io import BytesIO
+from typing import TYPE_CHECKING, Optional, Union
+
+if TYPE_CHECKING:
+    from .core import AudioModelFamilyV1
+
+logger = logging.getLogger(__name__)
+
+
+class F5TTSModel:
+    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._vocoder = None
+        self._kwargs = kwargs
+
+    @property
+    def model_ability(self):
+        return self._model_spec.model_ability
+
+    def load(self):
+        import os
+        import sys
+
+        # 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 f5_tts.infer.utils_infer import load_model, load_vocoder
+        from f5_tts.model import DiT
+
+        vocoder_name = self._kwargs.get("vocoder_name", "vocos")
+        vocoder_path = self._kwargs.get("vocoder_path")
+
+        if vocoder_name not in ["vocos", "bigvgan"]:
+            raise Exception(f"Unsupported vocoder name: {vocoder_name}")
+
+        if vocoder_path is not None:
+            self._vocoder = load_vocoder(
+                vocoder_name=vocoder_name, is_local=True, local_path=vocoder_path
+            )
+        else:
+            self._vocoder = load_vocoder(vocoder_name=vocoder_name, is_local=False)
+
+        model_cls = DiT
+        model_cfg = dict(
+            dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+        )
+        if vocoder_name == "vocos":
+            exp_name = "F5TTS_Base"
+            ckpt_step = 1200000
+        elif vocoder_name == "bigvgan":
+            exp_name = "F5TTS_Base_bigvgan"
+            ckpt_step = 1250000
+        else:
+            assert False
+        ckpt_file = os.path.join(
+            self._model_path, exp_name, f"model_{ckpt_step}.safetensors"
+        )
+        logger.info(f"Loading %s...", ckpt_file)
+        self._model = load_model(
+            model_cls, model_cfg, ckpt_file, mel_spec_type=vocoder_name
+        )
+
+    def _infer(self, ref_audio, ref_text, text_gen, model_obj, mel_spec_type, speed):
+        import numpy as np
+        from f5_tts.infer.utils_infer import infer_process, preprocess_ref_audio_text
+
+        config = {}
+        main_voice = {"ref_audio": ref_audio, "ref_text": ref_text}
+        if "voices" not in config:
+            voices = {"main": main_voice}
+        else:
+            voices = config["voices"]
+            voices["main"] = main_voice
+        for voice in voices:
+            (
+                voices[voice]["ref_audio"],
+                voices[voice]["ref_text"],
+            ) = preprocess_ref_audio_text(
+                voices[voice]["ref_audio"], 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 = []
+        reg1 = r"(?=\[\w+\])"
+        chunks = re.split(reg1, text_gen)
+        reg2 = r"\[(\w+)\]"
+        for text in chunks:
+            if not text.strip():
+                continue
+            match = re.match(reg2, text)
+            if match:
+                voice = match[1]
+            else:
+                logger.info("No voice tag found, using main.")
+                voice = "main"
+            if voice not in voices:
+                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"]
+            logger.info(f"Voice: {voice}")
+            audio, final_sample_rate, spectragram = infer_process(
+                ref_audio,
+                ref_text,
+                gen_text,
+                model_obj,
+                self._vocoder,
+                mel_spec_type=mel_spec_type,
+                speed=speed,
+            )
+            generated_audio_segments.append(audio)
+
+        if generated_audio_segments:
+            final_wave = np.concatenate(generated_audio_segments)
+            return final_sample_rate, final_wave
+        return None, None
+
+    def speech(
+        self,
+        input: str,
+        voice: str,
+        response_format: str = "mp3",
+        speed: float = 1.0,
+        stream: bool = False,
+        **kwargs,
+    ):
+        import f5_tts
+        import soundfile
+        import tomli
+
+        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)
+
+        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)
+            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=ref_audio,
+            ref_text=prompt_text,
+            text_gen=input,
+            model_obj=self._model,
+            mel_spec_type=vocoder_name,
+            speed=speed,
+        )
+
+        # Save the generated audio
+        with BytesIO() as out:
+            with soundfile.SoundFile(
+                out, "w", sample_rate, 1, format=response_format.upper()
+            ) as f:
+                f.write(wav)
+            return out.getvalue()

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()