diff --git a/modelscope/hub/snapshot_download.py b/modelscope/hub/snapshot_download.py index 75bcb991..8923e9e3 100644 --- a/modelscope/hub/snapshot_download.py +++ b/modelscope/hub/snapshot_download.py @@ -238,10 +238,8 @@ def _snapshot_download( headers = { 'user-agent': ModelScopeConfig.get_user_agent(user_agent=user_agent, ), + 'snapshot-identifier': str(uuid.uuid4()), } - if 'CI_TEST' not in os.environ: - # To count the download statistics, to add the snapshot-identifier as a header. - headers['snapshot-identifier'] = str(uuid.uuid4()) _api = HubApi() if cookies is None: cookies = ModelScopeConfig.get_cookies() @@ -258,12 +256,12 @@ def _snapshot_download( repo_id, revision=revision, cookies=cookies) revision = revision_detail['Revision'] - snapshot_header = headers if 'CI_TEST' in os.environ else { - **headers, - **{ - 'Snapshot': 'True' - } - } + # Add snapshot-ci-test for counting the ci test download + if 'CI_TEST' in os.environ: + snapshot_header = {**headers, **{'snapshot-ci-test': 'True'}} + else: + snapshot_header = {**headers, **{'Snapshot': 'True'}} + if cache.cached_model_revision is not None: snapshot_header[ 'cached_model_revision'] = cache.cached_model_revision diff --git a/modelscope/hub/utils/utils.py b/modelscope/hub/utils/utils.py index 7d377013..0fd078b0 100644 --- a/modelscope/hub/utils/utils.py +++ b/modelscope/hub/utils/utils.py @@ -106,8 +106,9 @@ def get_release_datetime(): def get_endpoint(): - modelscope_domain = os.getenv('MODELSCOPE_DOMAIN', - DEFAULT_MODELSCOPE_DOMAIN) + modelscope_domain = os.getenv( + 'MODELSCOPE_DOMAIN', + DEFAULT_MODELSCOPE_DOMAIN) or DEFAULT_MODELSCOPE_DOMAIN return MODELSCOPE_URL_SCHEME + modelscope_domain diff --git a/modelscope/utils/repo_utils.py b/modelscope/utils/repo_utils.py index 55d16251..d14adbf8 100644 --- a/modelscope/utils/repo_utils.py +++ b/modelscope/utils/repo_utils.py @@ -49,6 +49,7 @@ DATASET_LFS_SUFFIX = [ '.jack', '.jpeg', '.jpg', + '.png', '.jsonl', '.joblib', '.lz4',