[diffusion] chore: further improve model searching logic (#14484)

This commit is contained in:
Mick
2025-12-05 15:04:55 +08:00
committed by GitHub
parent b76e303e6a
commit beec8eed6a
2 changed files with 16 additions and 7 deletions

View File

@@ -11,7 +11,6 @@ import dataclasses
import importlib
import os
import pkgutil
import re
from functools import lru_cache
from typing import Any, Callable, Dict, List, Optional, Tuple, Type
@@ -158,6 +157,13 @@ def register_configs(
_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]:
"""
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)
# 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)
for model_hf_path in all_model_hf_paths:
if model_hf_path.lower() in cleaned_model_path:
for registered_model_hf_id in all_model_hf_paths:
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
if registered_model_name == model_name:
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)
# 3. Use detectors

View File

@@ -120,7 +120,8 @@ class ComponentLoader(ABC):
self.device = device
def should_offload(self, server_args, model_config: ModelConfig | None = None):
raise NotImplementedError()
# offload by default
return True
def target_device(self, should_offload):
if should_offload: