xinference 0.8.4__py3-none-any.whl → 0.9.0__py3-none-any.whl

xinference/model/embedding/core.py

@@ -13,27 +13,20 @@
 # limitations under the License.
 
 import logging
-import os
-import shutil
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple, Union, no_type_check
 
 import numpy as np
-from pydantic import BaseModel
 
-from ...constants import XINFERENCE_CACHE_DIR
 from ...types import Embedding, EmbeddingData, EmbeddingUsage
-from ..core import ModelDescription
-from ..utils import get_cache_dir, is_model_cached, valid_model_revision
+from ..core import CacheableModelSpec, ModelDescription
+from ..utils import get_cache_dir, is_model_cached
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_SCHEMES = ["s3"]
 # Used for check whether the model is cached.
 # Init when registering all the builtin models.
 MODEL_NAME_TO_REVISION: Dict[str, List[str]] = defaultdict(list)
-
-
 EMBEDDING_MODEL_DESCRIPTIONS: Dict[str, List[Dict]] = defaultdict(list)
 
 
@@ -43,7 +36,7 @@ def get_embedding_model_descriptions():
     return copy.deepcopy(EMBEDDING_MODEL_DESCRIPTIONS)
 
 
-class EmbeddingModelSpec(BaseModel):
+class EmbeddingModelSpec(CacheableModelSpec):
     model_name: str
     dimensions: int
     max_tokens: int
@@ -105,134 +98,10 @@ def generate_embedding_description(
     return res
 
 
-def cache_from_uri(
-    model_spec: EmbeddingModelSpec,
-    self_hosted_storage: bool = False,
-) -> str:
-    from fsspec import AbstractFileSystem, filesystem
-
-    from ..utils import copy_from_src_to_dst, parse_uri
-
-    cache_dir = os.path.realpath(
-        os.path.join(XINFERENCE_CACHE_DIR, model_spec.model_name)
-    )
-    if os.path.exists(cache_dir):
-        logger.info(f"Embedding cache {cache_dir} exists")
-        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.makedirs(XINFERENCE_CACHE_DIR, exist_ok=True)
-        os.symlink(src_root, cache_dir, target_is_directory=True)
-        return cache_dir
-    elif src_scheme in SUPPORTED_SCHEMES:
-        # use anonymous connection for self-hosted storage.
-        src_fs: AbstractFileSystem = filesystem(src_scheme, anon=self_hosted_storage)
-        local_fs: AbstractFileSystem = filesystem("file")
-
-        files_to_download = []
-        os.makedirs(cache_dir, exist_ok=True)
-
-        for path, _, files in src_fs.walk(model_spec.model_uri):
-            for file in files:
-                src_path = f"{path}/{file}"
-                local_path = src_path.replace(src_root, cache_dir)
-                files_to_download.append((src_path, local_path))
-
-        from concurrent.futures import ThreadPoolExecutor
-
-        failed = False
-        with ThreadPoolExecutor(max_workers=min(len(files_to_download), 4)) as executor:
-            futures = [
-                (
-                    src_path,
-                    executor.submit(
-                        copy_from_src_to_dst, src_fs, src_path, local_fs, local_path
-                    ),
-                )
-                for src_path, local_path in files_to_download
-            ]
-            for src_path, future in futures:
-                if failed:
-                    future.cancel()
-                else:
-                    try:
-                        future.result()
-                    except:
-                        logger.error(f"Download {src_path} failed", exc_info=True)
-                        failed = True
-
-        if failed:
-            logger.warning(f"Removing cache directory: {cache_dir}")
-            shutil.rmtree(cache_dir, ignore_errors=True)
-            raise RuntimeError(
-                f"Failed to download embedding model '{model_spec.model_name}' "
-            )
-        return cache_dir
-    else:
-        raise ValueError(f"Unsupported URL scheme: {src_scheme}")
-
-
 def cache(model_spec: EmbeddingModelSpec):
-    from huggingface_hub import snapshot_download as hf_download
-    from modelscope.hub.snapshot_download import snapshot_download as ms_download
-
-    from ..utils import retry_download, symlink_local_file
-
-    if (
-        hasattr(model_spec, "model_uri")
-        and getattr(model_spec, "model_uri", None) is not None
-    ):
-        logger.info(f"Embedding model caching from URI: {model_spec.model_uri}")
-        return cache_from_uri(model_spec=model_spec)
-
-    cache_dir = os.path.realpath(
-        os.path.join(XINFERENCE_CACHE_DIR, model_spec.model_name)
-    )
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-    meta_path = os.path.join(cache_dir, "__valid_download")
-    if valid_model_revision(meta_path, model_spec.model_revision):
-        return cache_dir
-
-    from_modelscope: bool = model_spec.model_hub == "modelscope"
-    if from_modelscope:
-        download_dir = retry_download(
-            ms_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-        )
-        for subdir, dirs, files in os.walk(download_dir):
-            for file in files:
-                relpath = os.path.relpath(os.path.join(subdir, file), download_dir)
-                symlink_local_file(os.path.join(subdir, file), cache_dir, relpath)
-    else:
-        retry_download(
-            hf_download,
-            model_spec.model_name,
-            None,
-            model_spec.model_id,
-            revision=model_spec.model_revision,
-            local_dir=cache_dir,
-            local_dir_use_symlinks=True,
-        )
-    with open(meta_path, "w") as f:
-        import json
+    from ..utils import cache
 
-        desc = EmbeddingModelDescription(None, None, model_spec)
-        json.dump(desc.to_dict(), f)
-    return cache_dir
+    return cache(model_spec, EmbeddingModelDescription)
 
 
 def get_cache_status(

xinference/model/image/__init__.py

@@ -15,6 +15,7 @@
 import codecs
 import json
 import os
+from itertools import chain
 
 from .core import (
     IMAGE_MODEL_DESCRIPTIONS,
@@ -25,13 +26,24 @@ from .core import (
 )
 
 _model_spec_json = os.path.join(os.path.dirname(__file__), "model_spec.json")
+_model_spec_modelscope_json = os.path.join(
+    os.path.dirname(__file__), "model_spec_modelscope.json"
+)
 BUILTIN_IMAGE_MODELS = dict(
     (spec["model_name"], ImageModelFamilyV1(**spec))
     for spec in json.load(codecs.open(_model_spec_json, "r", encoding="utf-8"))
 )
+MODELSCOPE_IMAGE_MODELS = dict(
+    (spec["model_name"], ImageModelFamilyV1(**spec))
+    for spec in json.load(
+        codecs.open(_model_spec_modelscope_json, "r", encoding="utf-8")
+    )
+)
 
 # register model description
-for model_name, model_spec in BUILTIN_IMAGE_MODELS.items():
+for model_name, model_spec in chain(
+    MODELSCOPE_IMAGE_MODELS.items(), BUILTIN_IMAGE_MODELS.items()
+):
     IMAGE_MODEL_DESCRIPTIONS.update(generate_image_description(model_spec))
 
 del _model_spec_json

xinference/model/image/core.py

@@ -17,10 +17,8 @@ import os
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
 
-from pydantic import BaseModel
-
 from ...constants import XINFERENCE_CACHE_DIR
-from ..core import ModelDescription
+from ..core import CacheableModelSpec, ModelDescription
 from ..utils import valid_model_revision
 from .stable_diffusion.core import DiffusionModel
 
@@ -37,11 +35,12 @@ def get_image_model_descriptions():
     return copy.deepcopy(IMAGE_MODEL_DESCRIPTIONS)
 
 
-class ImageModelFamilyV1(BaseModel):
+class ImageModelFamilyV1(CacheableModelSpec):
     model_family: str
     model_name: str
     model_id: str
     model_revision: str
+    model_hub: str = "huggingface"
     controlnet: Optional[List["ImageModelFamilyV1"]]
 
 
@@ -57,6 +56,10 @@ class ImageModelDescription(ModelDescription):
         self._model_spec = model_spec
 
     def to_dict(self):
+        if self._model_spec.controlnet is not None:
+            controlnet = [cn.dict() for cn in self._model_spec.controlnet]
+        else:
+            controlnet = self._model_spec.controlnet
         return {
             "model_type": "image",
             "address": self.address,
@@ -64,7 +67,7 @@ class ImageModelDescription(ModelDescription):
             "model_name": self._model_spec.model_name,
             "model_family": self._model_spec.model_family,
             "model_revision": self._model_spec.model_revision,
-            "controlnet": self._model_spec.controlnet,
+            "controlnet": controlnet,
         }
 
     def to_version_info(self):
@@ -111,7 +114,18 @@ def generate_image_description(
 
 
 def match_diffusion(model_name: str) -> ImageModelFamilyV1:
-    from . import BUILTIN_IMAGE_MODELS
+    from ..utils import download_from_modelscope
+    from . import BUILTIN_IMAGE_MODELS, MODELSCOPE_IMAGE_MODELS
+
+    if download_from_modelscope():
+        if model_name in MODELSCOPE_IMAGE_MODELS:
+            logger.debug(f"Image model {model_name} found in ModelScope.")
+            return MODELSCOPE_IMAGE_MODELS[model_name]
+        else:
+            logger.debug(
+                f"Image model {model_name} not found in ModelScope, "
+                f"now try to load it via builtin way."
+            )
 
     if model_name in BUILTIN_IMAGE_MODELS:
         return BUILTIN_IMAGE_MODELS[model_name]
@@ -123,44 +137,9 @@ def match_diffusion(model_name: str) -> ImageModelFamilyV1:
 
 
 def cache(model_spec: ImageModelFamilyV1):
-    # TODO: cache from uri
-    import huggingface_hub
-
-    cache_dir = get_cache_dir(model_spec)
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-
-    meta_path = os.path.join(cache_dir, "__valid_download")
-    if valid_model_revision(meta_path, model_spec.model_revision):
-        return cache_dir
-
-    for current_attempt in range(1, MAX_ATTEMPTS + 1):
-        try:
-            huggingface_hub.snapshot_download(
-                model_spec.model_id,
-                revision=model_spec.model_revision,
-                local_dir=cache_dir,
-                local_dir_use_symlinks=True,
-                resume_download=True,
-            )
-            break
-        except huggingface_hub.utils.LocalEntryNotFoundError:
-            remaining_attempts = MAX_ATTEMPTS - current_attempt
-            logger.warning(
-                f"Attempt {current_attempt} failed. Remaining attempts: {remaining_attempts}"
-            )
-    else:
-        raise RuntimeError(
-            f"Failed to download model '{model_spec.model_name}' after {MAX_ATTEMPTS} attempts"
-        )
-
-    with open(meta_path, "w") as f:
-        import json
-
-        desc = ImageModelDescription(None, None, model_spec)
-        json.dump(desc.to_dict(), f)
+    from ..utils import cache
 
-    return cache_dir
+    return cache(model_spec, ImageModelDescription)
 
 
 def get_cache_dir(model_spec: ImageModelFamilyV1):

xinference/model/image/model_spec_modelscope.json

@@ -0,0 +1,94 @@
+[
+  {
+    "model_name": "sd-turbo",
+    "model_family": "stable_diffusion",
+    "model_hub": "modelscope",
+    "model_id": "AI-ModelScope/sd-turbo",
+    "model_revision": "master"
+  },
+  {
+    "model_name": "sdxl-turbo",
+    "model_family": "stable_diffusion",
+    "model_hub": "modelscope",
+    "model_id": "AI-ModelScope/sdxl-turbo",
+    "model_revision": "master"
+  },
+  {
+    "model_name": "stable-diffusion-v1.5",
+    "model_family": "stable_diffusion",
+    "model_hub": "modelscope",
+    "model_id": "AI-ModelScope/stable-diffusion-v1-5",
+    "model_revision": "master",
+    "controlnet": [
+      {
+        "model_name":"canny",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-canny",
+        "model_revision": "7f2f69197050967007f6bbd23ab5e52f0384162a"
+      },
+      {
+        "model_name":"mlsd",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-mlsd",
+        "model_revision": "456fc1af4dc8f7d5fa6c949c5040d3828a24e3ad"
+      },
+      {
+        "model_name":"hed",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-hed",
+        "model_revision": "04473d9334ab44908daa66107bbfb6f710aa056d"
+      },
+      {
+        "model_name":"scribble",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-scribble",
+        "model_revision": "864edcd5ccc6ee2695eeebea5b4512100c83e7b3"
+      },
+      {
+        "model_name":"openpose",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-openpose",
+        "model_revision": "df796456519d1ac5dbe674caa0653fcc9673bca8"
+      },
+      {
+        "model_name":"normal",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-normal",
+        "model_revision": "1cbed9b3ca84422e4a2f23c14b9f5a114742b31d"
+      },
+      {
+        "model_name":"seg",
+        "model_family": "controlnet",
+        "model_id": "lllyasviel/sd-controlnet-seg",
+        "model_revision": "ecdcb5645b5099c9a7500a504fb9ab3f743c4d96"
+      }
+    ]
+  },
+  {
+    "model_name": "stable-diffusion-xl-base-1.0",
+    "model_family": "stable_diffusion",
+    "model_hub": "modelscope",
+    "model_id": "AI-ModelScope/stable-diffusion-xl-base-1.0",
+    "model_revision": "master",
+    "controlnet": [
+      {
+        "model_name":"canny",
+        "model_family": "controlnet",
+        "model_id": "diffusers/controlnet-canny-sdxl-1.0",
+        "model_revision": "eb115a19a10d14909256db740ed109532ab1483c"
+      },
+      {
+        "model_name":"depth",
+        "model_family": "controlnet",
+        "model_id": "diffusers/controlnet-depth-sdxl-1.0",
+        "model_revision": "cb2e660019b26f7d2379e26922b5b2c87253dcf7"
+      },
+      {
+        "model_name":"zoe-depth",
+        "model_family": "controlnet",
+        "model_id": "diffusers/controlnet-zoe-depth-sdxl-1.0",
+        "model_revision": "62134b9d8e703b5d6f74f1534457287a8bba77ef"
+      }
+    ]
+  }
+]

xinference/model/image/stable_diffusion/core.py

@@ -24,6 +24,7 @@ from io import BytesIO
 from typing import List, Optional, Union
 
 from ....constants import XINFERENCE_IMAGE_DIR
+from ....device_utils import move_model_to_available_device
 from ....types import Image, ImageList
 
 logger = logging.getLogger(__name__)
@@ -40,7 +41,7 @@ class DiffusionModel:
         self._kwargs = kwargs
 
     def load(self):
-        import torch
+        # import torch
         from diffusers import AutoPipelineForText2Image
 
         controlnet = self._kwargs.get("controlnet")
@@ -57,10 +58,7 @@ class DiffusionModel:
             # torch_dtype=torch.float16,
             # use_safetensors=True,
         )
-        if torch.cuda.is_available():
-            self._model = self._model.to("cuda")
-        elif torch.backends.mps.is_available():
-            self._model = self._model.to("mps")
+        self._model = move_model_to_available_device(self._model)
         # Recommended if your computer has < 64 GB of RAM
         self._model.enable_attention_slicing()
 

xinference/model/llm/ggml/llamacpp.py

@@ -307,7 +307,7 @@ class LlamaCppChatModel(LlamaCppModel, ChatModelMixin):
         generate_config = self._sanitize_generate_config(generate_config)
         # TODO(codingl2k1): qwen hacky to set stop for function call.
         model_family = self.model_family.model_family or self.model_family.model_name
-        if tools and "qwen-chat" == model_family:
+        if tools and model_family in ["qwen-chat", "qwen1.5-chat"]:
             stop = generate_config.get("stop")
             if isinstance(stop, str):
                 generate_config["stop"] = [stop, "Observation:"]