xinference 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

xinference/model/audio/core.py

@@ -14,7 +14,7 @@
 import logging
 import os
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple, Union
 
 from ...constants import XINFERENCE_CACHE_DIR
 from ..core import CacheableModelSpec, ModelDescription
@@ -94,7 +94,10 @@ def generate_audio_description(
     return res
 
 
-def match_audio(model_name: str) -> AudioModelFamilyV1:
+def match_audio(
+    model_name: str,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+) -> AudioModelFamilyV1:
     from ..utils import download_from_modelscope
     from . import BUILTIN_AUDIO_MODELS, MODELSCOPE_AUDIO_MODELS
     from .custom import get_user_defined_audios
@@ -103,17 +106,17 @@ def match_audio(model_name: str) -> AudioModelFamilyV1:
         if model_spec.model_name == model_name:
             return model_spec
 
-    if download_from_modelscope():
-        if model_name in MODELSCOPE_AUDIO_MODELS:
-            logger.debug(f"Audio model {model_name} found in ModelScope.")
-            return MODELSCOPE_AUDIO_MODELS[model_name]
-        else:
-            logger.debug(
-                f"Audio model {model_name} not found in ModelScope, "
-                f"now try to load it via builtin way."
-            )
-
-    if model_name in BUILTIN_AUDIO_MODELS:
+    if download_hub == "huggingface" and model_name in BUILTIN_AUDIO_MODELS:
+        logger.debug(f"Audio model {model_name} found in huggingface.")
+        return BUILTIN_AUDIO_MODELS[model_name]
+    elif download_hub == "modelscope" and model_name in MODELSCOPE_AUDIO_MODELS:
+        logger.debug(f"Audio model {model_name} found in ModelScope.")
+        return MODELSCOPE_AUDIO_MODELS[model_name]
+    elif download_from_modelscope() and model_name in MODELSCOPE_AUDIO_MODELS:
+        logger.debug(f"Audio model {model_name} found in ModelScope.")
+        return MODELSCOPE_AUDIO_MODELS[model_name]
+    elif model_name in BUILTIN_AUDIO_MODELS:
+        logger.debug(f"Audio model {model_name} found in huggingface.")
         return BUILTIN_AUDIO_MODELS[model_name]
     else:
         raise ValueError(
@@ -141,9 +144,14 @@ def get_cache_status(
 
 
 def create_audio_model_instance(
-    subpool_addr: str, devices: List[str], model_uid: str, model_name: str, **kwargs
+    subpool_addr: str,
+    devices: List[str],
+    model_uid: str,
+    model_name: str,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+    **kwargs,
 ) -> Tuple[Union[WhisperModel, ChatTTSModel], AudioModelDescription]:
-    model_spec = match_audio(model_name)
+    model_spec = match_audio(model_name, download_hub)
     model_path = cache(model_spec)
     model: Union[WhisperModel, ChatTTSModel]
     if model_spec.model_family == "whisper":

xinference/model/core.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Literal, Optional, Tuple, Union
 
 from .._compat import BaseModel
 from ..types import PeftModelConfig
@@ -55,10 +55,12 @@ def create_model_instance(
     model_size_in_billions: Optional[Union[int, str]] = None,
     quantization: Optional[str] = None,
     peft_model_config: Optional[PeftModelConfig] = None,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
     **kwargs,
 ) -> Tuple[Any, ModelDescription]:
     from .audio.core import create_audio_model_instance
     from .embedding.core import create_embedding_model_instance
+    from .flexible.core import create_flexible_model_instance
     from .image.core import create_image_model_instance
     from .llm.core import create_llm_model_instance
     from .rerank.core import create_rerank_model_instance
@@ -74,13 +76,14 @@ def create_model_instance(
             model_size_in_billions,
             quantization,
             peft_model_config,
+            download_hub,
             **kwargs,
         )
     elif model_type == "embedding":
         # 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, **kwargs
+            subpool_addr, devices, model_uid, model_name, download_hub, **kwargs
         )
     elif model_type == "image":
         kwargs.pop("trust_remote_code", None)
@@ -90,16 +93,22 @@ def create_model_instance(
             model_uid,
             model_name,
             peft_model_config,
+            download_hub,
             **kwargs,
         )
     elif model_type == "rerank":
         kwargs.pop("trust_remote_code", None)
         return create_rerank_model_instance(
-            subpool_addr, devices, model_uid, model_name, **kwargs
+            subpool_addr, devices, model_uid, model_name, download_hub, **kwargs
         )
     elif model_type == "audio":
         kwargs.pop("trust_remote_code", None)
         return create_audio_model_instance(
+            subpool_addr, devices, model_uid, model_name, download_hub, **kwargs
+        )
+    elif model_type == "flexible":
+        kwargs.pop("trust_remote_code", None)
+        return create_flexible_model_instance(
             subpool_addr, devices, model_uid, model_name, **kwargs
         )
     else:

xinference/model/embedding/core.py

@@ -16,7 +16,7 @@ import gc
 import logging
 import os
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple, Union, no_type_check
+from typing import Dict, List, Literal, Optional, Tuple, Union, no_type_check
 
 import numpy as np
 
@@ -305,7 +305,10 @@ class EmbeddingModel:
         )
 
 
-def match_embedding(model_name: str) -> EmbeddingModelSpec:
+def match_embedding(
+    model_name: str,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+) -> EmbeddingModelSpec:
     from ..utils import download_from_modelscope
     from . import BUILTIN_EMBEDDING_MODELS, MODELSCOPE_EMBEDDING_MODELS
     from .custom import get_user_defined_embeddings
@@ -315,29 +318,35 @@ def match_embedding(model_name: str) -> EmbeddingModelSpec:
         if model_name == model_spec.model_name:
             return model_spec
 
-    if download_from_modelscope():
-        if model_name in MODELSCOPE_EMBEDDING_MODELS:
-            logger.debug(f"Embedding model {model_name} found in ModelScope.")
-            return MODELSCOPE_EMBEDDING_MODELS[model_name]
-        else:
-            logger.debug(
-                f"Embedding model {model_name} not found in ModelScope, "
-                f"now try to load it via builtin way."
-            )
-
-    if model_name in BUILTIN_EMBEDDING_MODELS:
+    if download_hub == "modelscope" and model_name in MODELSCOPE_EMBEDDING_MODELS:
+        logger.debug(f"Embedding model {model_name} found in ModelScope.")
+        return MODELSCOPE_EMBEDDING_MODELS[model_name]
+    elif download_hub == "huggingface" and model_name in BUILTIN_EMBEDDING_MODELS:
+        logger.debug(f"Embedding model {model_name} found in Huggingface.")
+        return BUILTIN_EMBEDDING_MODELS[model_name]
+    elif download_from_modelscope() and model_name in MODELSCOPE_EMBEDDING_MODELS:
+        logger.debug(f"Embedding model {model_name} found in ModelScope.")
+        return MODELSCOPE_EMBEDDING_MODELS[model_name]
+    elif model_name in BUILTIN_EMBEDDING_MODELS:
+        logger.debug(f"Embedding model {model_name} found in Huggingface.")
         return BUILTIN_EMBEDDING_MODELS[model_name]
     else:
         raise ValueError(
             f"Embedding model {model_name} not found, available"
-            f"model list: {BUILTIN_EMBEDDING_MODELS.keys()}"
+            f"Huggingface: {BUILTIN_EMBEDDING_MODELS.keys()}"
+            f"ModelScope: {MODELSCOPE_EMBEDDING_MODELS.keys()}"
         )
 
 
 def create_embedding_model_instance(
-    subpool_addr: str, devices: List[str], model_uid: str, model_name: str, **kwargs
+    subpool_addr: str,
+    devices: List[str],
+    model_uid: str,
+    model_name: str,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+    **kwargs,
 ) -> Tuple[EmbeddingModel, EmbeddingModelDescription]:
-    model_spec = match_embedding(model_name)
+    model_spec = match_embedding(model_name, download_hub)
     model_path = cache(model_spec)
     model = EmbeddingModel(model_uid, model_path, **kwargs)
     model_description = EmbeddingModelDescription(

xinference/model/flexible/__init__.py

@@ -0,0 +1,40 @@
+# Copyright 2022-2024 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 os
+
+from ...constants import XINFERENCE_MODEL_DIR
+from .core import (
+    FLEXIBLE_MODEL_DESCRIPTIONS,
+    FlexibleModel,
+    FlexibleModelSpec,
+    generate_flexible_model_description,
+    get_flexible_model_descriptions,
+    get_flexible_models,
+    register_flexible_model,
+    unregister_flexible_model,
+)
+
+model_dir = os.path.join(XINFERENCE_MODEL_DIR, "flexible")
+if os.path.isdir(model_dir):
+    for f in os.listdir(model_dir):
+        with codecs.open(os.path.join(model_dir, f), encoding="utf-8") as fd:
+            model_spec = FlexibleModelSpec.parse_obj(json.load(fd))
+            register_flexible_model(model_spec, persist=False)
+
+# register model description
+for model in get_flexible_models():
+    FLEXIBLE_MODEL_DESCRIPTIONS.update(generate_flexible_model_description(model))

xinference/model/flexible/core.py

@@ -0,0 +1,228 @@
+# Copyright 2022-2024 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 json
+import logging
+import os
+from collections import defaultdict
+from threading import Lock
+from typing import Dict, List, Optional, Tuple
+
+from ...constants import XINFERENCE_CACHE_DIR, XINFERENCE_MODEL_DIR
+from ..core import CacheableModelSpec, ModelDescription
+from .utils import get_launcher
+
+logger = logging.getLogger(__name__)
+
+FLEXIBLE_MODEL_LOCK = Lock()
+
+
+class FlexibleModelSpec(CacheableModelSpec):
+    model_id: Optional[str]  # type: ignore
+    model_description: Optional[str]
+    model_uri: Optional[str]
+    launcher: str
+    launcher_args: Optional[str]
+
+    def parser_args(self):
+        return json.loads(self.launcher_args)
+
+
+class FlexibleModelDescription(ModelDescription):
+    def __init__(
+        self,
+        address: Optional[str],
+        devices: Optional[List[str]],
+        model_spec: FlexibleModelSpec,
+        model_path: Optional[str] = None,
+    ):
+        super().__init__(address, devices, model_path=model_path)
+        self._model_spec = model_spec
+
+    def to_dict(self):
+        return {
+            "model_type": "flexible",
+            "address": self.address,
+            "accelerators": self.devices,
+            "model_name": self._model_spec.model_name,
+            "launcher": self._model_spec.launcher,
+            "launcher_args": self._model_spec.launcher_args,
+        }
+
+    def get_model_version(self) -> str:
+        return f"{self._model_spec.model_name}"
+
+    def to_version_info(self):
+        return {
+            "model_version": self.get_model_version(),
+            "cache_status": True,
+            "model_file_location": self._model_spec.model_uri,
+            "launcher": self._model_spec.launcher,
+            "launcher_args": self._model_spec.launcher_args,
+        }
+
+
+def generate_flexible_model_description(
+    model_spec: FlexibleModelSpec,
+) -> Dict[str, List[Dict]]:
+    res = defaultdict(list)
+    res[model_spec.model_name].append(
+        FlexibleModelDescription(None, None, model_spec).to_version_info()
+    )
+    return res
+
+
+FLEXIBLE_MODELS: List[FlexibleModelSpec] = []
+FLEXIBLE_MODEL_DESCRIPTIONS: Dict[str, List[Dict]] = defaultdict(list)
+
+
+def get_flexible_models():
+    with FLEXIBLE_MODEL_LOCK:
+        return FLEXIBLE_MODELS.copy()
+
+
+def get_flexible_model_descriptions():
+    import copy
+
+    return copy.deepcopy(FLEXIBLE_MODEL_DESCRIPTIONS)
+
+
+def register_flexible_model(model_spec: FlexibleModelSpec, persist: bool):
+    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}.")
+
+    if model_spec.launcher_args:
+        try:
+            model_spec.parser_args()
+        except Exception:
+            raise ValueError(f"Invalid model launcher args {model_spec.launcher_args}.")
+
+    with FLEXIBLE_MODEL_LOCK:
+        for model_name in [spec.model_name for spec in FLEXIBLE_MODELS]:
+            if model_spec.model_name == model_name:
+                raise ValueError(
+                    f"Model name conflicts with existing model {model_spec.model_name}"
+                )
+        FLEXIBLE_MODELS.append(model_spec)
+
+    if persist:
+        persist_path = os.path.join(
+            XINFERENCE_MODEL_DIR, "flexible", f"{model_spec.model_name}.json"
+        )
+        os.makedirs(os.path.dirname(persist_path), exist_ok=True)
+        with open(persist_path, mode="w") as fd:
+            fd.write(model_spec.json())
+
+
+def unregister_flexible_model(model_name: str, raise_error: bool = True):
+    with FLEXIBLE_MODEL_LOCK:
+        model_spec = None
+        for i, f in enumerate(FLEXIBLE_MODELS):
+            if f.model_name == model_name:
+                model_spec = f
+                break
+        if model_spec:
+            FLEXIBLE_MODELS.remove(model_spec)
+
+            persist_path = os.path.join(
+                XINFERENCE_MODEL_DIR, "flexible", f"{model_spec.model_name}.json"
+            )
+            if os.path.exists(persist_path):
+                os.remove(persist_path)
+
+            cache_dir = os.path.join(XINFERENCE_CACHE_DIR, model_spec.model_name)
+            if os.path.exists(cache_dir):
+                logger.warning(
+                    f"Remove the cache of user-defined model {model_spec.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."
+                    )
+        else:
+            if raise_error:
+                raise ValueError(f"Model {model_name} not found")
+            else:
+                logger.warning(f"Model {model_name} not found")
+
+
+class FlexibleModel:
+    def __init__(
+        self,
+        model_uid: str,
+        model_path: str,
+        device: Optional[str] = None,
+        config: Optional[Dict] = None,
+    ):
+        self._model_uid = model_uid
+        self._model_path = model_path
+        self._device = device
+        self._config = config
+
+    def load(self):
+        """
+        Load the model.
+        """
+
+    def infer(self, **kwargs):
+        """
+        Call model to inference.
+        """
+        raise NotImplementedError("infer method not implemented.")
+
+    @property
+    def model_uid(self):
+        return self._model_uid
+
+    @property
+    def model_path(self):
+        return self._model_path
+
+    @property
+    def device(self):
+        return self._device
+
+    @property
+    def config(self):
+        return self._config
+
+
+def match_flexible_model(model_name):
+    for model_spec in get_flexible_models():
+        if model_name == model_spec.model_name:
+            return model_spec
+
+
+def create_flexible_model_instance(
+    subpool_addr: str, devices: List[str], model_uid: str, model_name: str, **kwargs
+) -> Tuple[FlexibleModel, FlexibleModelDescription]:
+    model_spec = match_flexible_model(model_name)
+    model_path = model_spec.model_uri
+    launcher_name = model_spec.launcher
+    launcher_args = model_spec.parser_args()
+    kwargs.update(launcher_args)
+
+    model = get_launcher(launcher_name)(
+        model_uid=model_uid, model_spec=model_spec, **kwargs
+    )
+
+    model_description = FlexibleModelDescription(
+        subpool_addr, devices, model_spec, model_path=model_path
+    )
+    return model, model_description

xinference/model/flexible/launchers/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2022-2024 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.
+
+from .transformers_launcher import launcher as transformers

xinference/model/flexible/launchers/transformers_launcher.py

@@ -0,0 +1,63 @@
+# Copyright 2022-2024 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.
+
+from transformers import pipeline
+
+from ..core import FlexibleModel, FlexibleModelSpec
+
+
+class MockModel(FlexibleModel):
+    def infer(self, **kwargs):
+        return kwargs
+
+
+class AutoModel(FlexibleModel):
+    def load(self):
+        config = self.config or {}
+        self._pipeline = pipeline(model=self.model_path, device=self.device, **config)
+
+    def infer(self, **kwargs):
+        return self._pipeline(**kwargs)
+
+
+class TransformersTextClassificationModel(FlexibleModel):
+    def load(self):
+        config = self.config or {}
+
+        self._pipeline = pipeline(model=self._model_path, device=self._device, **config)
+
+    def infer(self, **kwargs):
+        return self._pipeline(**kwargs)
+
+
+def launcher(model_uid: str, model_spec: FlexibleModelSpec, **kwargs) -> FlexibleModel:
+    task = kwargs.get("task")
+    device = kwargs.get("device")
+
+    model_path = model_spec.model_uri
+    if model_path is None:
+        raise ValueError("model_path required")
+
+    if task == "text-classification":
+        return TransformersTextClassificationModel(
+            model_uid=model_uid, model_path=model_path, device=device, config=kwargs
+        )
+    elif task == "mock":
+        return MockModel(
+            model_uid=model_uid, model_path=model_path, device=device, config=kwargs
+        )
+    else:
+        return AutoModel(
+            model_uid=model_uid, model_path=model_path, device=device, config=kwargs
+        )

xinference/model/flexible/utils.py

@@ -0,0 +1,33 @@
+# Copyright 2022-2024 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 importlib
+
+
+def get_launcher(launcher_name: str):
+    try:
+        i = launcher_name.rfind(".")
+        if i != -1:
+            module = importlib.import_module(launcher_name[:i])
+            fn = getattr(module, launcher_name[i + 1 :])
+        else:
+            importlib.import_module(launcher_name)
+            fn = locals().get(launcher_name)
+
+        if fn is None:
+            raise ValueError(f"Launcher {launcher_name} not found.")
+
+        return fn
+    except ImportError as e:
+        raise ImportError(f"Failed to import {launcher_name}: {e}")

xinference/model/image/core.py

@@ -15,7 +15,7 @@ import collections.abc
 import logging
 import os
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Literal, Optional, Tuple
 
 from ...constants import XINFERENCE_CACHE_DIR
 from ...types import PeftModelConfig
@@ -45,6 +45,7 @@ class ImageModelFamilyV1(CacheableModelSpec):
     model_id: str
     model_revision: str
     model_hub: str = "huggingface"
+    ability: Optional[str]
     controlnet: Optional[List["ImageModelFamilyV1"]]
 
 
@@ -71,6 +72,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,
+            "ability": self._model_spec.ability,
             "controlnet": controlnet,
         }
 
@@ -117,7 +119,10 @@ def generate_image_description(
     return res
 
 
-def match_diffusion(model_name: str) -> ImageModelFamilyV1:
+def match_diffusion(
+    model_name: str,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
+) -> ImageModelFamilyV1:
     from ..utils import download_from_modelscope
     from . import BUILTIN_IMAGE_MODELS, MODELSCOPE_IMAGE_MODELS
     from .custom import get_user_defined_images
@@ -126,17 +131,17 @@ def match_diffusion(model_name: str) -> ImageModelFamilyV1:
         if model_spec.model_name == model_name:
             return model_spec
 
-    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:
+    if download_hub == "modelscope" and model_name in MODELSCOPE_IMAGE_MODELS:
+        logger.debug(f"Image model {model_name} found in ModelScope.")
+        return MODELSCOPE_IMAGE_MODELS[model_name]
+    elif download_hub == "huggingface" and model_name in BUILTIN_IMAGE_MODELS:
+        logger.debug(f"Image model {model_name} found in Huggingface.")
+        return BUILTIN_IMAGE_MODELS[model_name]
+    elif download_from_modelscope() and model_name in MODELSCOPE_IMAGE_MODELS:
+        logger.debug(f"Image model {model_name} found in ModelScope.")
+        return MODELSCOPE_IMAGE_MODELS[model_name]
+    elif model_name in BUILTIN_IMAGE_MODELS:
+        logger.debug(f"Image model {model_name} found in Huggingface.")
         return BUILTIN_IMAGE_MODELS[model_name]
     else:
         raise ValueError(
@@ -183,9 +188,10 @@ def create_image_model_instance(
     model_uid: str,
     model_name: str,
     peft_model_config: Optional[PeftModelConfig] = None,
+    download_hub: Optional[Literal["huggingface", "modelscope", "csghub"]] = None,
     **kwargs,
 ) -> Tuple[DiffusionModel, ImageModelDescription]:
-    model_spec = match_diffusion(model_name)
+    model_spec = match_diffusion(model_name, download_hub)
     controlnet = kwargs.get("controlnet")
     # Handle controlnet
     if controlnet is not None:
@@ -230,6 +236,7 @@ def create_image_model_instance(
         lora_model_paths=lora_model,
         lora_load_kwargs=lora_load_kwargs,
         lora_fuse_kwargs=lora_fuse_kwargs,
+        ability=model_spec.ability,
         **kwargs,
     )
     model_description = ImageModelDescription(

xinference/model/image/custom.py

@@ -66,7 +66,7 @@ def register_image(model_spec: CustomImageModelFamilyV1, persist: bool):
             raise ValueError(f"Invalid model URI {model_uri}")
 
         persist_path = os.path.join(
-            XINFERENCE_MODEL_DIR, "image", f"{model_spec.model_id}.json"
+            XINFERENCE_MODEL_DIR, "image", f"{model_spec.model_name}.json"
         )
         os.makedirs(os.path.dirname(persist_path), exist_ok=True)
         with open(persist_path, "w") as f: