support downloading resource

This commit is contained in:
Artiprocher
2025-11-18 10:29:07 +08:00
parent 8332ecebb7
commit 1313f4dd63
2 changed files with 20 additions and 8 deletions

View File

@@ -4,6 +4,7 @@ from PIL import Image
from einops import repeat, reduce
from typing import Optional, Union
from dataclasses import dataclass
from huggingface_hub import snapshot_download as hf_snapshot_download
from modelscope import snapshot_download
import numpy as np
from PIL import Image
@@ -196,13 +197,24 @@ class ModelConfig:
self.local_model_path = "./models"
if not skip_download:
downloaded_files = glob.glob(self.origin_file_pattern, root_dir=os.path.join(self.local_model_path, self.model_id))
snapshot_download(
self.model_id,
local_dir=os.path.join(self.local_model_path, self.model_id),
allow_file_pattern=allow_file_pattern,
ignore_file_pattern=downloaded_files,
local_files_only=False
)
if self.download_resource.lower() == "modelscope":
snapshot_download(
self.model_id,
local_dir=os.path.join(self.local_model_path, self.model_id),
allow_file_pattern=allow_file_pattern,
ignore_file_pattern=downloaded_files,
local_files_only=False
)
elif self.download_resource.lower() == "huggingface":
hf_snapshot_download(
self.model_id,
local_dir=os.path.join(self.local_model_path, self.model_id),
allow_patterns=allow_file_pattern,
ignore_patterns=downloaded_files,
local_files_only=False
)
else:
raise ValueError("`download_resource` should be `modelscope` or `huggingface`.")
# Let rank 1, 2, ... wait for rank 0
if use_usp: