[diffusion] feat: add ModelScope support (#17924)

This commit is contained in:
yrk111222
2026-02-09 19:23:45 +08:00
committed by GitHub
parent 615a02dcd4
commit 76eb1c8406
3 changed files with 94 additions and 17 deletions

View File

@@ -8,7 +8,8 @@ from functools import lru_cache
from typing import Optional
import filelock
from huggingface_hub import hf_hub_download
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
@@ -43,6 +44,8 @@ def _maybe_download_model(
Local directory path that contains the downloaded config file, or the original local directory.
"""
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import hf_hub_download
if os.path.exists(model_name_or_path):
logger.info("Model already exists locally")
return model_name_or_path
@@ -52,9 +55,11 @@ def _maybe_download_model(
with _get_lock(model_name_or_path):
# Try `model_index.json` first (diffusers models)
source_hub = "MS Hub" if envs.SGLANG_USE_MODELSCOPE.get() else "HF Hub"
try:
logger.info(
"Downloading model_index.json from HF Hub for %s...",
"Downloading model_index.json from %s for %s...",
source_hub,
model_name_or_path,
)
file_path = hf_hub_download(
@@ -70,7 +75,9 @@ def _maybe_download_model(
# Fallback to `config.json`
try:
logger.info(
"Downloading config.json from HF Hub for %s...", model_name_or_path
"Downloading config.json from %s for %s...",
source_hub,
model_name_or_path,
)
file_path = hf_hub_download(
repo_id=model_name_or_path,
@@ -83,9 +90,9 @@ def _maybe_download_model(
raise ValueError(
(
"Could not find model locally at %s and failed to download "
"model_index.json/config.json from HF Hub: %s"
"model_index.json/config.json from %s: %s"
)
% (model_name_or_path, e_config)
% (model_name_or_path, source_hub, e_config)
) from e_config

View File

@@ -12,6 +12,7 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
from sglang.multimodal_gen.runtime.scheduler_client import sync_scheduler_client
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.environ import envs
logger = init_logger(__name__)
@@ -27,13 +28,21 @@ def run_sgl_diffusion_webui(server_args: ServerArgs):
# import gradio in function to avoid CI crash
import gradio as gr
from huggingface_hub import model_info
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope.hub.api import HubApi
api = HubApi()
model_info_obj = api.model_info(server_args.model_path)
task_name = model_info_obj.tasks[0]["Name"].replace("-synthesis", "")
else:
from huggingface_hub import model_info
task_name = model_info(server_args.model_path).pipeline_tag
# init client
sync_scheduler_client.initialize(server_args)
task_name = model_info(server_args.model_path).pipeline_tag
if task_name in ("text-to-video", "image-to-video", "video-to-video"):
task_type = "video"
elif task_name in ["text-to-image", "image-to-image"]:

View File

@@ -26,12 +26,11 @@ import shutil
import time
from functools import reduce
from pathlib import Path
from typing import Any, Optional, cast
from typing import Any, Optional, Union, cast
from diffusers.loaders.lora_base import (
_best_guess_weight_name, # watch out for potetential removal from diffusers
)
from huggingface_hub import snapshot_download
from huggingface_hub.errors import (
LocalEntryNotFoundError,
RepositoryNotFoundError,
@@ -45,6 +44,7 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_N
from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.environ import envs
from sglang.utils import is_in_ci
logger = init_logger(__name__)
@@ -467,7 +467,6 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
"""
import tempfile
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import EntryNotFoundError
# If it's a local path, verify it directly
@@ -638,9 +637,7 @@ def maybe_download_model(
ignore_patterns=["*.onnx", "*.msgpack"],
local_dir=local_dir,
local_files_only=True,
resume_download=True,
max_workers=8,
etag_timeout=60,
)
if is_lora or _verify_model_complete(local_path):
# CI validation: check all subdirectories for missing shards
@@ -709,9 +706,7 @@ def maybe_download_model(
ignore_patterns=["*.onnx", "*.msgpack"],
allow_patterns=allow_patterns,
local_dir=local_dir,
resume_download=True,
max_workers=8,
etag_timeout=120,
)
# Verify downloaded model is complete (skip for LoRA)
@@ -725,9 +720,7 @@ def maybe_download_model(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],
local_dir=local_dir,
resume_download=True,
max_workers=8,
etag_timeout=60,
force_download=True,
)
if not _verify_model_complete(local_path):
@@ -775,3 +768,71 @@ def maybe_download_model(
raise ValueError(
f"Could not find model at {model_name_or_path} and failed to download from HF Hub: {e}"
) from e
# Unified download functions with Hugging Face-compatible names
def hf_hub_download(
repo_id: str,
filename: str,
local_dir: Optional[Union[str, Path]] = None,
**kwargs,
) -> str:
"""Unified hf_hub_download that supports both Hugging Face Hub and ModelScope."""
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import model_file_download
return model_file_download(
model_id=repo_id,
file_path=filename,
cache_dir=local_dir,
**kwargs,
)
else:
from huggingface_hub import hf_hub_download as _hf_hub_download
return _hf_hub_download(
repo_id=repo_id,
filename=filename,
local_dir=local_dir,
**kwargs,
)
def snapshot_download(
repo_id: str,
local_dir: Optional[Union[str, Path]] = None,
ignore_patterns: Optional[Union[list[str], str]] = None,
allow_patterns: Optional[Union[list[str], str]] = None,
local_files_only: bool = False,
max_workers: int = 8,
**kwargs,
) -> str:
"""Unified snapshot_download that supports both Hugging Face Hub and ModelScope."""
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import snapshot_download as _ms_snapshot_download
ms_kwargs = {
"model_id": repo_id,
"local_dir": local_dir,
"ignore_patterns": ignore_patterns,
"allow_patterns": allow_patterns,
"local_files_only": local_files_only,
"max_workers": max_workers,
}
ms_kwargs.update(kwargs)
return _ms_snapshot_download(**ms_kwargs)
else:
from huggingface_hub import snapshot_download as _hf_snapshot_download
hf_kwargs = {
"repo_id": repo_id,
"local_dir": local_dir,
"ignore_patterns": ignore_patterns,
"allow_patterns": allow_patterns,
"local_files_only": local_files_only,
"max_workers": max_workers,
"resume_download": True,
"etag_timeout": 60,
}
hf_kwargs.update(kwargs)
return _hf_snapshot_download(**hf_kwargs)