114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
"""
|
||
|
|
Standalone model download script.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python model_download.py
|
||
|
|
"""
|
||
|
|
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from dotenv import load_dotenv
|
||
|
|
|
||
|
|
DEFAULT_ENV_FILE = ".env"
|
||
|
|
DEFAULT_MODEL_ID = "Qwen/Qwen3.5-9B"
|
||
|
|
DEFAULT_MODEL_DIR = "./models/Qwen3.5-9B"
|
||
|
|
DEFAULT_CACHE_DIR = "./modelscope_cache"
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_path(raw: str, base_dir: Path) -> Path:
|
||
|
|
path = Path(raw).expanduser()
|
||
|
|
if path.is_absolute():
|
||
|
|
return path.resolve()
|
||
|
|
return (base_dir / path).resolve()
|
||
|
|
|
||
|
|
|
||
|
|
def load_runtime_env(script_dir: Path) -> Path:
|
||
|
|
env_name = (os.getenv("ENV_FILE", DEFAULT_ENV_FILE) or DEFAULT_ENV_FILE).strip()
|
||
|
|
env_path = (script_dir / env_name).resolve()
|
||
|
|
if env_path.exists():
|
||
|
|
load_dotenv(env_path)
|
||
|
|
return env_path
|
||
|
|
|
||
|
|
|
||
|
|
def env_flag(name: str, default: bool = False) -> bool:
|
||
|
|
raw = os.getenv(name, "")
|
||
|
|
if not raw:
|
||
|
|
return default
|
||
|
|
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||
|
|
|
||
|
|
|
||
|
|
def download_model(
|
||
|
|
model_id: str,
|
||
|
|
model_dir: Path,
|
||
|
|
cache_dir: Path,
|
||
|
|
revision: str = "",
|
||
|
|
skip_if_exists: bool = False,
|
||
|
|
) -> Path:
|
||
|
|
try:
|
||
|
|
from modelscope.hub.snapshot_download import snapshot_download
|
||
|
|
except Exception as exc:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Missing dependencies. Please install first:\n"
|
||
|
|
" pip install -r requirements.txt"
|
||
|
|
) from exc
|
||
|
|
|
||
|
|
model_dir.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
if skip_if_exists and model_dir.exists() and any(model_dir.iterdir()):
|
||
|
|
print(f"[INFO] model already exists, skip download: {model_dir}")
|
||
|
|
return model_dir.resolve()
|
||
|
|
|
||
|
|
print(f"[INFO] model_id={model_id}")
|
||
|
|
print(f"[INFO] model_dir={model_dir}")
|
||
|
|
print(f"[INFO] cache_dir={cache_dir}")
|
||
|
|
if revision:
|
||
|
|
print(f"[INFO] revision={revision}")
|
||
|
|
|
||
|
|
kwargs = {
|
||
|
|
"model_id": model_id,
|
||
|
|
"local_dir": str(model_dir),
|
||
|
|
"cache_dir": str(cache_dir),
|
||
|
|
}
|
||
|
|
if revision:
|
||
|
|
kwargs["revision"] = revision
|
||
|
|
|
||
|
|
downloaded_path = snapshot_download(**kwargs)
|
||
|
|
print(f"[OK] download complete: {downloaded_path}")
|
||
|
|
return Path(downloaded_path).resolve()
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
script_dir = Path(__file__).resolve().parent
|
||
|
|
env_path = load_runtime_env(script_dir)
|
||
|
|
model_id = os.getenv("DOWNLOAD_MODEL_ID", os.getenv("MODEL_SOURCE", DEFAULT_MODEL_ID)).strip()
|
||
|
|
model_dir_raw = os.getenv("DOWNLOAD_SAVE_DIR", os.getenv("MODEL_DIR", DEFAULT_MODEL_DIR)).strip()
|
||
|
|
cache_dir_raw = os.getenv("DOWNLOAD_CACHE_DIR", os.getenv("MODELSCOPE_CACHE", DEFAULT_CACHE_DIR)).strip()
|
||
|
|
revision = os.getenv("DOWNLOAD_REVISION", "").strip()
|
||
|
|
skip_if_exists = env_flag("SKIP_MODEL_DOWNLOAD_IF_EXISTS", True)
|
||
|
|
|
||
|
|
if not model_id:
|
||
|
|
raise ValueError("DOWNLOAD_MODEL_ID/MODEL_SOURCE is empty.")
|
||
|
|
if not model_dir_raw:
|
||
|
|
raise ValueError("DOWNLOAD_SAVE_DIR/MODEL_DIR is empty.")
|
||
|
|
if not cache_dir_raw:
|
||
|
|
raise ValueError("DOWNLOAD_CACHE_DIR/MODELSCOPE_CACHE is empty.")
|
||
|
|
|
||
|
|
model_dir = resolve_path(model_dir_raw, script_dir)
|
||
|
|
cache_dir = resolve_path(cache_dir_raw, script_dir)
|
||
|
|
print(f"[INFO] env_file={env_path}")
|
||
|
|
download_model(
|
||
|
|
model_id=model_id,
|
||
|
|
model_dir=model_dir,
|
||
|
|
cache_dir=cache_dir,
|
||
|
|
revision=revision,
|
||
|
|
skip_if_exists=skip_if_exists,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|