xinference 1.7.1.post1__py3-none-any.whl → 1.8.0__py3-none-any.whl

xinference/model/audio/utils.py

@@ -21,15 +21,9 @@ from collections.abc import Callable
 import numpy as np
 import torch
 
-from .core import AudioModelFamilyV1
-
 logger = logging.getLogger(__name__)
 
 
-def get_model_version(audio_model: AudioModelFamilyV1) -> str:
-    return audio_model.model_name
-
-
 def _extract_pcm_from_wav_bytes(wav_bytes):
     with io.BytesIO(wav_bytes) as wav_io:
         with wave.open(wav_io, "rb") as wav_file:

xinference/model/audio/whisper.py

@@ -26,7 +26,7 @@ from ...device_utils import (
 )
 
 if TYPE_CHECKING:
-    from .core import AudioModelFamilyV1
+    from .core import AudioModelFamilyV2
 
 logger = logging.getLogger(__name__)
 
@@ -43,11 +43,12 @@ class WhisperModel:
         self,
         model_uid: str,
         model_path: str,
-        model_spec: "AudioModelFamilyV1",
+        model_spec: "AudioModelFamilyV2",
         device: Optional[str] = None,
         max_new_tokens: Optional[int] = 128,
         **kwargs,
     ):
+        self.model_family = model_spec
         self._model_uid = model_uid
         self._model_path = model_path
         self._model_spec = model_spec

xinference/model/audio/whisper_mlx.py

@@ -18,7 +18,7 @@ import tempfile
 from typing import TYPE_CHECKING, List, Optional
 
 if TYPE_CHECKING:
-    from .core import AudioModelFamilyV1
+    from .core import AudioModelFamilyV2
 
 logger = logging.getLogger(__name__)
 
@@ -28,10 +28,11 @@ class WhisperMLXModel:
         self,
         model_uid: str,
         model_path: str,
-        model_spec: "AudioModelFamilyV1",
+        model_spec: "AudioModelFamilyV2",
         device: Optional[str] = None,
         **kwargs,
     ):
+        self.model_family = model_spec
         self._model_uid = model_uid
         self._model_path = model_path
         self._model_spec = model_spec

xinference/model/cache_manager.py

@@ -0,0 +1,141 @@
+import logging
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .core import CacheableModelSpec
+
+
+logger = logging.getLogger(__name__)
+
+
+class CacheManager:
+    def __init__(self, model_family: "CacheableModelSpec"):
+        from ..constants import XINFERENCE_CACHE_DIR, XINFERENCE_MODEL_DIR
+
+        self._model_family = model_family
+        self._v2_cache_dir_prefix = os.path.join(XINFERENCE_CACHE_DIR, "v2")
+        self._v2_custom_dir_prefix = os.path.join(XINFERENCE_MODEL_DIR, "v2")
+        os.makedirs(self._v2_cache_dir_prefix, exist_ok=True)
+        os.makedirs(self._v2_custom_dir_prefix, exist_ok=True)
+        self._cache_dir = os.path.join(
+            self._v2_cache_dir_prefix, self._model_family.model_name.replace(".", "_")
+        )
+
+    def get_cache_dir(self):
+        return self._cache_dir
+
+    def get_cache_status(self):
+        cache_dir = self.get_cache_dir()
+        return os.path.exists(cache_dir)
+
+    def _cache_from_uri(self, model_spec: "CacheableModelSpec") -> str:
+        from .utils import parse_uri
+
+        cache_dir = self.get_cache_dir()
+        if os.path.exists(cache_dir):
+            logger.info("cache %s exists", cache_dir)
+            return cache_dir
+
+        assert model_spec.model_uri is not None
+        src_scheme, src_root = parse_uri(model_spec.model_uri)
+        if src_root.endswith("/"):
+            # remove trailing path separator.
+            src_root = src_root[:-1]
+
+        if src_scheme == "file":
+            if not os.path.isabs(src_root):
+                raise ValueError(
+                    f"Model URI cannot be a relative path: {model_spec.model_uri}"
+                )
+            os.symlink(src_root, cache_dir, target_is_directory=True)
+            return cache_dir
+        else:
+            raise ValueError(f"Unsupported URL scheme: {src_scheme}")
+
+    def _cache(self) -> str:
+        from .utils import IS_NEW_HUGGINGFACE_HUB, create_symlink, retry_download
+
+        if (
+            hasattr(self._model_family, "model_uri")
+            and getattr(self._model_family, "model_uri", None) is not None
+        ):
+            logger.info(f"Model caching from URI: {self._model_family.model_uri}")
+            return self._cache_from_uri(model_spec=self._model_family)
+
+        cache_dir = self.get_cache_dir()
+        if self.get_cache_status():
+            return cache_dir
+
+        from_modelscope: bool = self._model_family.model_hub == "modelscope"
+        cache_config = (
+            self._model_family.cache_config.copy()
+            if self._model_family.cache_config
+            else {}
+        )
+        if from_modelscope:
+            from modelscope.hub.snapshot_download import (
+                snapshot_download as ms_download,
+            )
+
+            download_dir = retry_download(
+                ms_download,
+                self._model_family.model_name,
+                None,
+                self._model_family.model_id,
+                revision=self._model_family.model_revision,
+                **cache_config,
+            )
+            create_symlink(download_dir, cache_dir)
+        else:
+            from huggingface_hub import snapshot_download as hf_download
+
+            use_symlinks = cache_config
+            if not IS_NEW_HUGGINGFACE_HUB:
+                use_symlinks = {"local_dir_use_symlinks": True, "local_dir": cache_dir}
+            download_dir = retry_download(
+                hf_download,
+                self._model_family.model_name,
+                None,
+                self._model_family.model_id,
+                revision=self._model_family.model_revision,
+                **use_symlinks,
+            )
+            if IS_NEW_HUGGINGFACE_HUB:
+                create_symlink(download_dir, cache_dir)
+        return cache_dir
+
+    def cache(self) -> str:
+        return self._cache()
+
+    def register_custom_model(self, model_type: str):
+        persist_path = os.path.join(
+            self._v2_custom_dir_prefix,
+            model_type,
+            f"{self._model_family.model_name}.json",
+        )
+        os.makedirs(os.path.dirname(persist_path), exist_ok=True)
+        with open(persist_path, mode="w") as fd:
+            fd.write(self._model_family.json())
+
+    def unregister_custom_model(self, model_type: str):
+        persist_path = os.path.join(
+            self._v2_custom_dir_prefix,
+            model_type,
+            f"{self._model_family.model_name}.json",
+        )
+        if os.path.exists(persist_path):
+            os.remove(persist_path)
+
+        cache_dir = self.get_cache_dir()
+        if self.get_cache_status():
+            logger.warning(
+                f"Remove the cache of user-defined model {self._model_family.model_name}. "
+                f"Cache directory: {cache_dir}"
+            )
+            if os.path.islink(cache_dir):
+                os.remove(cache_dir)
+            else:
+                logger.warning(
+                    f"Cache directory is not a soft link, please remove it manually."
+                )

xinference/model/core.py

@@ -11,47 +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.
-
-from abc import ABC, abstractmethod
-from typing import Any, List, Literal, Optional, Tuple, Union
+from typing import Any, List, Literal, Optional, Union
 
 from .._compat import BaseModel
 from ..types import PeftModelConfig
 
 
-class ModelDescription(ABC):
-    def __init__(
-        self,
-        address: Optional[str],
-        devices: Optional[List[str]],
-        model_path: Optional[str] = None,
-    ):
-        self.address = address
-        self.devices = devices
-        self._model_path = model_path
-
-    @property
-    @abstractmethod
-    def spec(self):
-        pass
-
-    def to_dict(self):
-        """
-        Return a dict to describe some information about model.
-        :return:
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def to_version_info(self):
-        """
-        Return a dict to describe version info about a model instance
-        """
-
-
 def create_model_instance(
-    subpool_addr: str,
-    devices: List[str],
     model_uid: str,
     model_type: str,
     model_name: str,
@@ -65,7 +31,7 @@ def create_model_instance(
     ] = None,
     model_path: Optional[str] = None,
     **kwargs,
-) -> Tuple[Any, ModelDescription]:
+) -> Any:
     from .audio.core import create_audio_model_instance
     from .embedding.core import create_embedding_model_instance
     from .flexible.core import create_flexible_model_instance
@@ -76,8 +42,6 @@ def create_model_instance(
 
     if model_type == "LLM":
         return create_llm_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             model_engine,
@@ -93,11 +57,11 @@ def create_model_instance(
         # embedding model doesn't accept trust_remote_code
         kwargs.pop("trust_remote_code", None)
         return create_embedding_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             model_engine,
+            model_format,
+            quantization,
             download_hub,
             model_path,
             **kwargs,
@@ -105,8 +69,6 @@ def create_model_instance(
     elif model_type == "image":
         kwargs.pop("trust_remote_code", None)
         return create_image_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             peft_model_config,
@@ -117,8 +79,6 @@ def create_model_instance(
     elif model_type == "rerank":
         kwargs.pop("trust_remote_code", None)
         return create_rerank_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             download_hub,
@@ -128,8 +88,6 @@ def create_model_instance(
     elif model_type == "audio":
         kwargs.pop("trust_remote_code", None)
         return create_audio_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             download_hub,
@@ -139,8 +97,6 @@ def create_model_instance(
     elif model_type == "video":
         kwargs.pop("trust_remote_code", None)
         return create_video_model_instance(
-            subpool_addr,
-            devices,
             model_uid,
             model_name,
             download_hub,
@@ -150,7 +106,7 @@ def create_model_instance(
     elif model_type == "flexible":
         kwargs.pop("trust_remote_code", None)
         return create_flexible_model_instance(
-            subpool_addr, devices, model_uid, model_name, model_path, **kwargs
+            model_uid, model_name, model_path, **kwargs
         )
     else:
         raise ValueError(f"Unsupported model type: {model_type}.")
@@ -161,6 +117,7 @@ class CacheableModelSpec(BaseModel):
     model_id: str
     model_revision: Optional[str]
     model_hub: str = "huggingface"
+    cache_config: Optional[dict]
 
 
 class VirtualEnvSettings(BaseModel):

xinference/model/custom.py

@@ -0,0 +1,174 @@
+# Copyright 2022-2025 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 codecs
+import json
+import logging
+import os
+import threading
+import warnings
+from typing import TYPE_CHECKING, Dict, List, Type
+
+if TYPE_CHECKING:
+    from .core import CacheableModelSpec
+
+logger = logging.getLogger(__name__)
+
+
+class ModelRegistry:
+    model_type = "unknown"
+
+    def __init__(self) -> None:
+        self.lock = threading.Lock()
+        self.models: List["CacheableModelSpec"] = []
+        self.builtin_models: List[str] = []
+
+    def find_model(self, model_name: str):
+        model_spec = None
+        for f in self.models:
+            if f.model_name == model_name:
+                model_spec = f
+                break
+        return model_spec
+
+    def get_custom_models(self):
+        with self.lock:
+            return self.models.copy()
+
+    def check_model_uri(self, model_spec: "CacheableModelSpec"):
+        from .utils import is_valid_model_uri
+
+        model_uri = model_spec.model_uri
+        if model_uri and not is_valid_model_uri(model_uri):
+            raise ValueError(f"Invalid model URI {model_uri}.")
+
+    def add_ud_model(self, model_spec):
+        self.models.append(model_spec)
+
+    def register(self, model_spec: "CacheableModelSpec", persist: bool):
+        from .cache_manager import CacheManager
+        from .utils import is_valid_model_name
+
+        if not is_valid_model_name(model_spec.model_name):
+            raise ValueError(f"Invalid model name {model_spec.model_name}.")
+
+        self.check_model_uri(model_spec)
+
+        with self.lock:
+            for model_name in self.builtin_models + [
+                spec.model_name for spec in self.models
+            ]:
+                if model_spec.model_name == model_name:
+                    raise ValueError(
+                        f"Model name conflicts with existing model {model_spec.model_name}"
+                    )
+
+            self.add_ud_model(model_spec)
+
+        if persist:
+            cache_manager = CacheManager(model_spec)
+            cache_manager.register_custom_model(self.model_type)
+
+    def remove_ud_model(self, model_spec):
+        self.models.remove(model_spec)
+
+    def remove_ud_model_files(self, model_spec):
+        from .cache_manager import CacheManager
+
+        cache_manager = CacheManager(model_spec)
+        cache_manager.unregister_custom_model(self.model_type)
+
+    def unregister(
+        self, model_name: str, raise_error: bool = True, remove_file: bool = True
+    ):
+        with self.lock:
+            model_spec = self.find_model(model_name)
+            if model_spec:
+                self.remove_ud_model(model_spec)
+                if remove_file:
+                    self.remove_ud_model_files(model_spec)
+            else:
+                if raise_error:
+                    raise ValueError(f"Model {model_name} not found")
+                else:
+                    logger.warning(
+                        f"Custom {self.model_type} model {model_name} not found"
+                    )
+
+
+class RegistryManager:
+    _instances: Dict[str, ModelRegistry] = {}
+
+    @classmethod
+    def get_registry(cls, model_type: str) -> ModelRegistry:
+        from .audio.custom import AudioModelRegistry
+        from .embedding.custom import EmbeddingModelRegistry
+        from .flexible.custom import FlexibleModelRegistry
+        from .image.custom import ImageModelRegistry
+        from .llm.custom import LLMModelRegistry
+        from .rerank.custom import RerankModelRegistry
+
+        if model_type not in cls._instances:
+            if model_type == "rerank":
+                cls._instances[model_type] = RerankModelRegistry()
+            elif model_type == "image":
+                cls._instances[model_type] = ImageModelRegistry()
+            elif model_type == "audio":
+                cls._instances[model_type] = AudioModelRegistry()
+            elif model_type == "llm":
+                cls._instances[model_type] = LLMModelRegistry()
+            elif model_type == "flexible":
+                cls._instances[model_type] = FlexibleModelRegistry()
+            elif model_type == "embedding":
+                cls._instances[model_type] = EmbeddingModelRegistry()
+            else:
+                raise ValueError(f"Unknown model type: {model_type}")
+        return cls._instances[model_type]
+
+
+def migrate_from_v1_to_v2(model_type: str, model_spec_cls: Type):
+    from ..constants import XINFERENCE_MODEL_DIR
+
+    v1_user_defined_model_dir = os.path.join(XINFERENCE_MODEL_DIR, model_type)
+    v2_user_defined_model_dir = os.path.join(XINFERENCE_MODEL_DIR, "v2", model_type)
+    if os.path.isdir(v1_user_defined_model_dir):
+        for f in os.listdir(v1_user_defined_model_dir):
+            if os.path.exists(os.path.join(v2_user_defined_model_dir, f)):
+                # skip if v2 has already
+                continue
+
+            try:
+                with codecs.open(
+                    os.path.join(v1_user_defined_model_dir, f), encoding="utf-8"
+                ) as fd:
+                    v1_model_json = json.load(fd)
+
+                    v1_model_json["version"] = 2
+                    for spec in v1_model_json.get("model_specs", []):
+                        if "quantizations" in spec:
+                            # change quantizations to quantization
+                            spec["quantization"] = spec["quantizations"][0]
+
+                    user_defined_model_family = model_spec_cls(**v1_model_json)
+                    registry = RegistryManager.get_registry(model_type)
+                    # register custom model file to v2
+                    registry.register(user_defined_model_family, persist=True)
+                    # unregister since it will be registered by v2
+                    registry.unregister(
+                        user_defined_model_family.model_name, remove_file=False
+                    )
+            except Exception as e:
+                warnings.warn(
+                    f"Fail to migrate {v1_user_defined_model_dir}/{f}, error: {e}"
+                )