[diffusion] chore: further improve model searching logic (#14484)
This commit is contained in:
@@ -11,7 +11,6 @@ import dataclasses
|
|||||||
import importlib
|
import importlib
|
||||||
import os
|
import os
|
||||||
import pkgutil
|
import pkgutil
|
||||||
import re
|
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
|
||||||
|
|
||||||
@@ -158,6 +157,13 @@ def register_configs(
|
|||||||
_MODEL_NAME_DETECTORS.append((model_id, detector))
|
_MODEL_NAME_DETECTORS.append((model_id, detector))
|
||||||
|
|
||||||
|
|
||||||
|
def get_model_short_name(model_id: str) -> str:
|
||||||
|
if "/" in model_id:
|
||||||
|
return model_id.split("/")[-1]
|
||||||
|
else:
|
||||||
|
return model_id
|
||||||
|
|
||||||
|
|
||||||
def _get_config_info(model_path: str) -> Optional[ConfigInfo]:
|
def _get_config_info(model_path: str) -> Optional[ConfigInfo]:
|
||||||
"""
|
"""
|
||||||
Gets the ConfigInfo for a given model path using mappings and detectors.
|
Gets the ConfigInfo for a given model path using mappings and detectors.
|
||||||
@@ -169,14 +175,16 @@ def _get_config_info(model_path: str) -> Optional[ConfigInfo]:
|
|||||||
return _CONFIG_REGISTRY.get(model_id)
|
return _CONFIG_REGISTRY.get(model_id)
|
||||||
|
|
||||||
# 2. Partial match: find the best (longest) match against all registered model hf paths.
|
# 2. Partial match: find the best (longest) match against all registered model hf paths.
|
||||||
cleaned_model_path = re.sub(r"--", "/", model_path.lower())
|
model_name = get_model_short_name(model_path.lower())
|
||||||
all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)
|
all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)
|
||||||
for model_hf_path in all_model_hf_paths:
|
for registered_model_hf_id in all_model_hf_paths:
|
||||||
if model_hf_path.lower() in cleaned_model_path:
|
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
|
||||||
|
|
||||||
|
if registered_model_name == model_name:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
f"Resolved model name '{model_hf_path}' from partial path match."
|
f"Resolved model name '{registered_model_hf_id}' from partial path match."
|
||||||
)
|
)
|
||||||
model_id = _MODEL_HF_PATH_TO_NAME[model_hf_path]
|
model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id]
|
||||||
return _CONFIG_REGISTRY.get(model_id)
|
return _CONFIG_REGISTRY.get(model_id)
|
||||||
|
|
||||||
# 3. Use detectors
|
# 3. Use detectors
|
||||||
|
|||||||
@@ -120,7 +120,8 @@ class ComponentLoader(ABC):
|
|||||||
self.device = device
|
self.device = device
|
||||||
|
|
||||||
def should_offload(self, server_args, model_config: ModelConfig | None = None):
|
def should_offload(self, server_args, model_config: ModelConfig | None = None):
|
||||||
raise NotImplementedError()
|
# offload by default
|
||||||
|
return True
|
||||||
|
|
||||||
def target_device(self, should_offload):
|
def target_device(self, should_offload):
|
||||||
if should_offload:
|
if should_offload:
|
||||||
|
|||||||
Reference in New Issue
Block a user