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

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,6 +81,11 @@ class FishSpeechModel:
             if not is_device_available(self._device):
                 raise ValueError(f"Device {self._device} is not available!")
 
+        # 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)
@@ -102,102 +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,
-        seed="0",
-        streaming=False,
-    ):
-        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,
-            GenerateResponse,
-            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,
-            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=self._kwargs.get("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,
-        )
-
-        response_queue = queue.Queue()
-        self._llama_queue.put(
-            GenerateRequest(
-                request=request,
-                response_queue=response_queue,
-            )
+        self._engine = TTSInferenceEngine(
+            self._llama_queue, self._model, precision, enable_compile
         )
 
-        segments = []
-
-        while True:
-            result: WrappedGenerateResponse = response_queue.get()
-            if result.status == "error":
-                raise result.response
-
-            result: GenerateResponse = result.response
-            if result.action == "next":
-                break
-
-            with autocast_exclude_mps(
-                device_type=self._model.device.type,
-                dtype=self._kwargs.get("precision", 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, 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,
@@ -211,22 +124,31 @@ class FishSpeechModel:
         if speed != 1.0:
             logger.warning("Fish speech does not support setting speed: %s.", speed)
         import torchaudio
+        from tools.schema import ServeReferenceAudio, ServeTTSRequest
 
         prompt_speech = kwargs.get("prompt_speech")
         prompt_text = kwargs.get("prompt_text", kwargs.get("reference_text", ""))
-        result = self._inference(
-            text=input,
-            enable_reference_audio=kwargs.get(
-                "enable_reference_audio", prompt_speech is not None
-            ),
-            reference_audio=prompt_speech,
-            reference_text=prompt_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,
+        if prompt_speech is not None:
+            r = ServeReferenceAudio(audio=prompt_speech, text=prompt_text)
+            references = [r]
+        else:
+            references = []
+
+        assert self._engine is not None
+        result = self._engine.inference(
+            ServeTTSRequest(
+                text=input,
+                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,
+            )
         )
 
         if stream:
@@ -242,7 +164,9 @@ class FishSpeechModel:
                     last_pos = 0
                     with writer.open():
                         for chunk in result:
-                            chunk = chunk[0]
+                            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)
@@ -257,7 +181,7 @@ class FishSpeechModel:
             return _stream_generator()
         else:
             result = list(result)
-            sample_rate, audio = result[0][1]
+            sample_rate, audio = result[0].audio
             audio = np.array([audio])
 
             # Save the generated audio

xinference/model/audio/model_spec.json

@@ -236,10 +236,18 @@
     "multilingual": true
   },
   {
-    "model_name": "FishSpeech-1.4",
+    "model_name": "CosyVoice2-0.5B",
+    "model_family": "CosyVoice",
+    "model_id": "mrfakename/CosyVoice2-0.5B",
+    "model_revision": "5676baabc8a76dc93ef60a88bbd2420deaa2f644",
+    "model_ability": "text-to-audio",
+    "multilingual": true
+  },
+  {
+    "model_name": "FishSpeech-1.5",
     "model_family": "FishAudio",
-    "model_id": "fishaudio/fish-speech-1.4",
-    "model_revision": "069c573759936b35191d3380deb89183c0656f59",
+    "model_id": "fishaudio/fish-speech-1.5",
+    "model_revision": "268b6ec86243dd683bc78dab7e9a6cedf9191f2a",
     "model_ability": "text-to-audio",
     "multilingual": true
   },
@@ -250,5 +258,13 @@
     "model_revision": "4dcc16f297f2ff98a17b3726b16f5de5a5e45672",
     "model_ability": "text-to-audio",
     "multilingual": true
+  },
+  {
+    "model_name": "F5-TTS-MLX",
+    "model_family": "F5-TTS-MLX",
+    "model_id": "lucasnewman/f5-tts-mlx",
+    "model_revision": "7642bb232e3fcacf92c51c786edebb8624da6b93",
+    "model_ability": "text-to-audio",
+    "multilingual": true
   }
 ]

xinference/model/audio/model_spec_modelscope.json

@@ -74,6 +74,15 @@
     "model_ability": "text-to-audio",
     "multilingual": true
   },
+  {
+    "model_name": "CosyVoice2-0.5B",
+    "model_family": "CosyVoice",
+    "model_hub": "modelscope",
+    "model_id": "iic/CosyVoice2-0.5B",
+    "model_revision": "master",
+    "model_ability": "text-to-audio",
+    "multilingual": true
+  },
   {
     "model_name": "F5-TTS",
     "model_family": "F5-TTS",

xinference/model/audio/utils.py

@@ -11,8 +11,40 @@
 # 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 numpy as np
+
 from .core import AudioModelFamilyV1
 
 
 def get_model_version(audio_model: AudioModelFamilyV1) -> str:
     return audio_model.model_name
+
+
+def ensure_sample_rate(
+    audio: np.ndarray, old_sample_rate: int, sample_rate: int
+) -> np.ndarray:
+    import soundfile as sf
+    from scipy.signal import resample
+
+    if old_sample_rate != sample_rate:
+        # Calculate the new data length
+        new_length = int(len(audio) * sample_rate / old_sample_rate)
+
+        # Resample the data
+        resampled_data = resample(audio, new_length)
+
+        # Use BytesIO to save the resampled data to memory
+        with io.BytesIO() as buffer:
+            # Write the resampled data to the memory buffer
+            sf.write(buffer, resampled_data, sample_rate, format="WAV")
+
+            # Reset the buffer position to the beginning
+            buffer.seek(0)
+
+            # Read the data from the memory buffer
+            audio, sr = sf.read(buffer, dtype="float32")
+
+    return audio

xinference/model/image/core.py

@@ -22,7 +22,12 @@ from typing import Dict, List, Literal, Optional, Tuple, Union
 from ...constants import XINFERENCE_CACHE_DIR
 from ...types import PeftModelConfig
 from ..core import CacheableModelSpec, ModelDescription
-from ..utils import valid_model_revision
+from ..utils import (
+    IS_NEW_HUGGINGFACE_HUB,
+    retry_download,
+    symlink_local_file,
+    valid_model_revision,
+)
 from .ocr.got_ocr2 import GotOCR2Model
 from .stable_diffusion.core import DiffusionModel
 from .stable_diffusion.mlx import MLXDiffusionModel
@@ -51,6 +56,9 @@ class ImageModelFamilyV1(CacheableModelSpec):
     controlnet: Optional[List["ImageModelFamilyV1"]]
     default_model_config: Optional[dict] = {}
     default_generate_config: Optional[dict] = {}
+    gguf_model_id: Optional[str]
+    gguf_quantizations: Optional[List[str]]
+    gguf_model_file_name_template: Optional[str]
 
 
 class ImageModelDescription(ModelDescription):
@@ -187,6 +195,61 @@ def get_cache_status(
         return valid_model_revision(meta_path, model_spec.model_revision)
 
 
+def cache_gguf(spec: ImageModelFamilyV1, quantization: Optional[str] = None):
+    if not quantization:
+        return
+
+    cache_dir = os.path.realpath(os.path.join(XINFERENCE_CACHE_DIR, spec.model_name))
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir, exist_ok=True)
+
+    if not spec.gguf_model_file_name_template:
+        raise NotImplementedError(
+            f"{spec.model_name} does not support GGUF quantization"
+        )
+    if quantization not in (spec.gguf_quantizations or []):
+        raise ValueError(
+            f"Cannot support quantization {quantization}, "
+            f"available quantizations: {spec.gguf_quantizations}"
+        )
+
+    filename = spec.gguf_model_file_name_template.format(quantization=quantization)  # type: ignore
+    full_path = os.path.join(cache_dir, filename)
+
+    if spec.model_hub == "huggingface":
+        import huggingface_hub
+
+        use_symlinks = {}
+        if not IS_NEW_HUGGINGFACE_HUB:
+            use_symlinks = {"local_dir_use_symlinks": True, "local_dir": cache_dir}
+        download_file_path = retry_download(
+            huggingface_hub.hf_hub_download,
+            spec.model_name,
+            None,
+            spec.gguf_model_id,
+            filename=filename,
+            **use_symlinks,
+        )
+        if IS_NEW_HUGGINGFACE_HUB:
+            symlink_local_file(download_file_path, cache_dir, filename)
+    elif spec.model_hub == "modelscope":
+        from modelscope.hub.file_download import model_file_download
+
+        download_file_path = retry_download(
+            model_file_download,
+            spec.model_name,
+            None,
+            spec.gguf_model_id,
+            filename,
+            revision=spec.model_revision,
+        )
+        symlink_local_file(download_file_path, cache_dir, filename)
+    else:
+        raise NotImplementedError
+
+    return full_path
+
+
 def create_ocr_model_instance(
     subpool_addr: str,
     devices: List[str],
@@ -219,6 +282,8 @@ def create_image_model_instance(
         Literal["huggingface", "modelscope", "openmind_hub", "csghub"]
     ] = None,
     model_path: Optional[str] = None,
+    gguf_quantization: Optional[str] = None,
+    gguf_model_path: Optional[str] = None,
     **kwargs,
 ) -> Tuple[
     Union[DiffusionModel, MLXDiffusionModel, GotOCR2Model], ImageModelDescription
@@ -272,6 +337,8 @@ def create_image_model_instance(
             ]
     if not model_path:
         model_path = cache(model_spec)
+    if not gguf_model_path and gguf_quantization:
+        gguf_model_path = cache_gguf(model_spec, gguf_quantization)
     if peft_model_config is not None:
         lora_model = peft_model_config.peft_model
         lora_load_kwargs = peft_model_config.image_lora_load_kwargs
@@ -298,6 +365,7 @@ def create_image_model_instance(
         lora_load_kwargs=lora_load_kwargs,
         lora_fuse_kwargs=lora_fuse_kwargs,
         model_spec=model_spec,
+        gguf_model_path=gguf_model_path,
         **kwargs,
     )
     model_description = ImageModelDescription(

xinference/model/image/model_spec.json

@@ -11,8 +11,24 @@
     ],
     "default_model_config": {
       "quantize": true,
-      "quantize_text_encoder": "text_encoder_2"
-    }
+      "quantize_text_encoder": "text_encoder_2",
+      "torch_dtype": "bfloat16"
+    },
+    "gguf_model_id": "city96/FLUX.1-schnell-gguf",
+    "gguf_quantizations": [
+      "F16",
+      "Q2_K",
+      "Q3_K_S",
+      "Q4_0",
+      "Q4_1",
+      "Q4_K_S",
+      "Q5_0",
+      "Q5_1",
+      "Q5_K_S",
+      "Q6_K",
+      "Q8_0"
+    ],
+    "gguf_model_file_name_template": "flux1-schnell-{quantization}.gguf"
   },
   {
     "model_name": "FLUX.1-dev",
@@ -26,8 +42,24 @@
     ],
     "default_model_config": {
       "quantize": true,
-      "quantize_text_encoder": "text_encoder_2"
-    }
+      "quantize_text_encoder": "text_encoder_2",
+      "torch_dtype": "bfloat16"
+    },
+    "gguf_model_id": "city96/FLUX.1-dev-gguf",
+    "gguf_quantizations": [
+      "F16",
+      "Q2_K",
+      "Q3_K_S",
+      "Q4_0",
+      "Q4_1",
+      "Q4_K_S",
+      "Q5_0",
+      "Q5_1",
+      "Q5_K_S",
+      "Q6_K",
+      "Q8_0"
+    ],
+    "gguf_model_file_name_template": "flux1-dev-{quantization}.gguf"
   },
   {
     "model_name": "sd3-medium",
@@ -44,6 +76,97 @@
       "quantize_text_encoder": "text_encoder_3"
     }
   },
+  {
+    "model_name": "sd3.5-medium",
+    "model_family": "stable_diffusion",
+    "model_id": "stabilityai/stable-diffusion-3.5-medium",
+    "model_revision": "94b13ccbe959c51e8159d91f562c58f29fac971a",
+    "model_ability": [
+      "text2image",
+      "image2image",
+      "inpainting"
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_3",
+      "torch_dtype": "bfloat16"
+    },
+    "gguf_model_id": "city96/stable-diffusion-3.5-medium-gguf",
+    "gguf_quantizations": [
+      "F16",
+      "Q3_K_M",
+      "Q3_K_S",
+      "Q4_0",
+      "Q4_1",
+      "Q4_K_M",
+      "Q4_K_S",
+      "Q5_0",
+      "Q5_1",
+      "Q5_K_M",
+      "Q5_K_S",
+      "Q6_K",
+      "Q8_0"
+    ],
+    "gguf_model_file_name_template": "sd3.5_medium-{quantization}.gguf"
+  },
+  {
+    "model_name": "sd3.5-large",
+    "model_family": "stable_diffusion",
+    "model_id": "stabilityai/stable-diffusion-3.5-large",
+    "model_revision": "ceddf0a7fdf2064ea28e2213e3b84e4afa170a0f",
+    "model_ability": [
+      "text2image",
+      "image2image",
+      "inpainting"
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_3",
+      "torch_dtype": "bfloat16",
+      "transformer_nf4": true
+    },
+    "gguf_model_id": "city96/stable-diffusion-3.5-large-gguf",
+    "gguf_quantizations": [
+      "F16",
+      "Q4_0",
+      "Q4_1",
+      "Q5_0",
+      "Q5_1",
+      "Q8_0"
+    ],
+    "gguf_model_file_name_template": "sd3.5_large-{quantization}.gguf"
+  },
+  {
+    "model_name": "sd3.5-large-turbo",
+    "model_family": "stable_diffusion",
+    "model_id": "stabilityai/stable-diffusion-3.5-large-turbo",
+    "model_revision": "ec07796fc06b096cc56de9762974a28f4c632eda",
+    "model_ability": [
+      "text2image",
+      "image2image",
+      "inpainting"
+    ],
+    "default_model_config": {
+      "quantize": true,
+      "quantize_text_encoder": "text_encoder_3",
+      "torch_dtype": "bfloat16",
+      "transformer_nf4": true
+    },
+    "default_generate_config": {
+      "guidance_scale": 1.0,
+      "num_inference_steps": 4
+    },
+    "gguf_model_id": "city96/stable-diffusion-3.5-large-turbo-gguf",
+    "gguf_quantizations": [
+      "F16",
+      "Q4_0",
+      "Q4_1",
+      "Q5_0",
+      "Q5_1",
+      "Q8_0"
+    ],
+    "gguf_model_file_name_template": "sd3.5_large_turbo-{quantization}.gguf"
+  },
   {
     "model_name": "sd-turbo",
     "model_family": "stable_diffusion",