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

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

xinference/model/audio/fish_speech.py

@@ -11,10 +11,8 @@
 # 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 gc
 import logging
 import os.path
-import queue
 import sys
 from io import BytesIO
 from typing import TYPE_CHECKING, Optional
@@ -60,6 +58,7 @@ class FishSpeechModel:
         self._device = device
         self._llama_queue = None
         self._model = None
+        self._engine = None
         self._kwargs = kwargs
 
     @property
@@ -72,6 +71,7 @@ class FishSpeechModel:
             0, os.path.join(os.path.dirname(__file__), "../../thirdparty/fish_speech")
         )
 
+        from tools.inference_engine import TTSInferenceEngine
         from tools.llama.generate import launch_thread_safe_queue
         from tools.vqgan.inference import load_model as load_decoder_model
 
@@ -81,12 +81,19 @@ class FishSpeechModel:
             if not is_device_available(self._device):
                 raise ValueError(f"Device {self._device} is not available!")
 
-        logger.info("Loading Llama model...")
+        # https://github.com/pytorch/pytorch/issues/129207
+        if self._device == "mps":
+            logger.warning("The Conv1d has bugs on MPS backend, fallback to CPU.")
+            self._device = "cpu"
+
+        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...")
 
@@ -100,98 +107,10 @@ class FishSpeechModel:
             device=self._device,
         )
 
-    @torch.inference_mode()
-    def _inference(
-        self,
-        text,
-        enable_reference_audio,
-        reference_audio,
-        reference_text,
-        max_new_tokens,
-        chunk_length,
-        top_p,
-        repetition_penalty,
-        temperature,
-        streaming=False,
-    ):
-        from fish_speech.utils import autocast_exclude_mps
-        from tools.api import decode_vq_tokens, encode_reference
-        from tools.llama.generate import (
-            GenerateRequest,
-            GenerateResponse,
-            WrappedGenerateResponse,
-        )
-
-        # Parse reference audio aka prompt
-        prompt_tokens = encode_reference(
-            decoder_model=self._model,
-            reference_audio=reference_audio,
-            enable_reference_audio=enable_reference_audio,
-        )
-
-        # LLAMA Inference
-        request = dict(
-            device=self._model.device,
-            max_new_tokens=max_new_tokens,
-            text=text,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            temperature=temperature,
-            compile=False,
-            iterative_prompt=chunk_length > 0,
-            chunk_length=chunk_length,
-            max_length=2048,
-            prompt_tokens=prompt_tokens if enable_reference_audio else None,
-            prompt_text=reference_text if enable_reference_audio else None,
+        self._engine = TTSInferenceEngine(
+            self._llama_queue, self._model, precision, enable_compile
         )
 
-        response_queue = queue.Queue()
-        self._llama_queue.put(
-            GenerateRequest(
-                request=request,
-                response_queue=response_queue,
-            )
-        )
-
-        if streaming:
-            yield wav_chunk_header(), None, None
-
-        segments = []
-
-        while True:
-            result: WrappedGenerateResponse = response_queue.get()  # type: ignore
-            if result.status == "error":
-                raise Exception(str(result.response))
-
-            result: GenerateResponse = result.response  # type: ignore
-            if result.action == "next":
-                break
-
-            with autocast_exclude_mps(
-                device_type=self._model.device.type, dtype=torch.bfloat16
-            ):
-                fake_audios = decode_vq_tokens(
-                    decoder_model=self._model,
-                    codes=result.codes,
-                )
-
-            fake_audios = fake_audios.float().cpu().numpy()
-            segments.append(fake_audios)
-
-            if streaming:
-                yield (fake_audios * 32768).astype(np.int16).tobytes(), None, None
-
-        if len(segments) == 0:
-            raise Exception("No audio generated, please check the input text.")
-
-        # No matter streaming or not, we need to return the final audio
-        audio = np.concatenate(segments, axis=0)
-        yield None, (self._model.spec_transform.sample_rate, audio), None
-
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
-
     def speech(
         self,
         input: str,
@@ -204,29 +123,70 @@ 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
+        from tools.schema import ServeReferenceAudio, ServeTTSRequest
+
+        prompt_speech = kwargs.get("prompt_speech")
+        prompt_text = kwargs.get("prompt_text", kwargs.get("reference_text", ""))
+        if prompt_speech is not None:
+            r = ServeReferenceAudio(audio=prompt_speech, text=prompt_text)
+            references = [r]
+        else:
+            references = []
 
-        result = list(
-            self._inference(
+        assert self._engine is not None
+        result = self._engine.inference(
+            ServeTTSRequest(
                 text=input,
-                enable_reference_audio=False,
-                reference_audio=None,
-                reference_text=kwargs.get("reference_text", ""),
+                references=references,
+                reference_id=kwargs.get("reference_id"),
+                seed=kwargs.get("seed"),
                 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,
+                format=response_format,
             )
         )
-        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:
+                            if chunk.code == "final":
+                                continue
+                            chunk = chunk.audio[1]
+                            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].audio
+            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()