153 lines
5.1 KiB
Python
153 lines
5.1 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from typing import Optional
|
|
|
|
import filelock
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
temp_dir = tempfile.gettempdir()
|
|
|
|
|
|
def _get_lock(model_name_or_path: str, cache_dir: Optional[str] = None):
|
|
lock_dir = cache_dir or temp_dir
|
|
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
|
model_name = model_name_or_path.replace("/", "-")
|
|
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
|
lock_file_name = hash_name + model_name + ".lock"
|
|
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
|
return lock
|
|
|
|
|
|
# Copied and adapted from hf_diffusers_utils.py
|
|
def _maybe_download_model(
|
|
model_name_or_path: str, local_dir: str | None = None, download: bool = True
|
|
) -> str:
|
|
"""
|
|
Resolve a model path. If it's a local directory, return it.
|
|
If it's a Hugging Face Hub ID, download only the config file
|
|
(`model_index.json` or `config.json`) and return its directory.
|
|
|
|
Args:
|
|
model_name_or_path: Local path or Hugging Face Hub model ID
|
|
local_dir: Local directory to save the downloaded file (if any)
|
|
download: Whether to download from Hugging Face Hub when needed
|
|
|
|
Returns:
|
|
Local directory path that contains the downloaded config file, or the original local directory.
|
|
"""
|
|
|
|
if os.path.exists(model_name_or_path):
|
|
logger.info("Model already exists locally")
|
|
return model_name_or_path
|
|
|
|
if not download:
|
|
return model_name_or_path
|
|
|
|
with _get_lock(model_name_or_path):
|
|
# Try `model_index.json` first (diffusers models)
|
|
try:
|
|
logger.info(
|
|
"Downloading model_index.json from HF Hub for %s...",
|
|
model_name_or_path,
|
|
)
|
|
file_path = hf_hub_download(
|
|
repo_id=model_name_or_path,
|
|
filename="model_index.json",
|
|
local_dir=local_dir,
|
|
)
|
|
logger.info("Downloaded to %s", file_path)
|
|
return os.path.dirname(file_path)
|
|
except Exception as e_index:
|
|
logger.debug("model_index.json not found or failed: %s", e_index)
|
|
|
|
# Fallback to `config.json`
|
|
try:
|
|
logger.info(
|
|
"Downloading config.json from HF Hub for %s...", model_name_or_path
|
|
)
|
|
file_path = hf_hub_download(
|
|
repo_id=model_name_or_path,
|
|
filename="config.json",
|
|
local_dir=local_dir,
|
|
)
|
|
logger.info("Downloaded to %s", file_path)
|
|
return os.path.dirname(file_path)
|
|
except Exception as e_config:
|
|
raise ValueError(
|
|
(
|
|
"Could not find model locally at %s and failed to download "
|
|
"model_index.json/config.json from HF Hub: %s"
|
|
)
|
|
% (model_name_or_path, e_config)
|
|
) from e_config
|
|
|
|
|
|
# Copied and adapted from hf_diffusers_utils.py
|
|
def is_diffusers_model_path(model_path: str) -> True:
|
|
"""
|
|
Verify if the model directory contains a valid diffusers configuration.
|
|
|
|
Args:
|
|
model_path: Path to the model directory
|
|
|
|
Returns:
|
|
The loaded model configuration as a dictionary if the model is a diffusers model
|
|
None if the model is not a diffusers model
|
|
"""
|
|
|
|
# Prefer model_index.json which indicates a diffusers pipeline
|
|
config_path = os.path.join(model_path, "model_index.json")
|
|
if not os.path.exists(config_path):
|
|
return False
|
|
|
|
# Load the config
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
# Verify diffusers version exists
|
|
if "_diffusers_version" not in config:
|
|
return False
|
|
return True
|
|
|
|
|
|
def get_is_diffusion_model(model_path: str):
|
|
model_path = _maybe_download_model(model_path)
|
|
is_diffusion_model = is_diffusers_model_path(model_path)
|
|
if is_diffusion_model:
|
|
logger.info("Diffusion model detected")
|
|
return is_diffusion_model
|
|
|
|
|
|
def get_model_path(extra_argv):
|
|
# Find the model_path argument
|
|
model_path = None
|
|
for i, arg in enumerate(extra_argv):
|
|
if arg == "--model-path":
|
|
if i + 1 < len(extra_argv):
|
|
model_path = extra_argv[i + 1]
|
|
break
|
|
elif arg.startswith("--model-path="):
|
|
model_path = arg.split("=", 1)[1]
|
|
break
|
|
|
|
if model_path is None:
|
|
# Fallback for --help or other cases where model-path is not provided
|
|
if any(h in extra_argv for h in ["-h", "--help"]):
|
|
raise Exception(
|
|
"Usage: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
|
|
"This command can launch either a standard language model server or a diffusion model server.\n"
|
|
"The server type is determined by the model path.\n"
|
|
"For specific arguments, please provide a model_path."
|
|
)
|
|
else:
|
|
raise Exception(
|
|
"Error: --model-path is required. "
|
|
"Please provide the path to the model."
|
|
)
|
|
return model_path
|