diff --git a/modelscope/cli/download.py b/modelscope/cli/download.py index 6b430453..42ad9062 100644 --- a/modelscope/cli/download.py +++ b/modelscope/cli/download.py @@ -3,12 +3,13 @@ import os from argparse import ArgumentParser from modelscope.cli.base import CLICommand +from modelscope.hub.api import HubApi, ModelScopeConfig from modelscope.hub.constants import DEFAULT_MAX_WORKERS from modelscope.hub.file_download import (dataset_file_download, model_file_download) from modelscope.hub.snapshot_download import (dataset_snapshot_download, snapshot_download) -from modelscope.hub.utils.utils import convert_patterns +from modelscope.hub.utils.utils import convert_patterns, get_endpoint from modelscope.utils.constant import DEFAULT_DATASET_REVISION @@ -54,16 +55,21 @@ class DownloadCMD(CLICommand): default='model', help="Type of repo to download from (defaults to 'model').", ) + parser.add_argument( + '--token', + type=str, + default=None, + help='Optional. Access token to download controlled entities.') parser.add_argument( '--revision', type=str, default=None, - help='Revision of the model.') + help='Revision of the entity (e.g., model).') parser.add_argument( '--cache_dir', type=str, default=None, - help='Cache directory to save model.') + help='Cache directory to save entity (e.g., model).') parser.add_argument( '--local_dir', type=str, @@ -118,6 +124,11 @@ class DownloadCMD(CLICommand): % self.args.repo_type) if not self.args.model and not self.args.dataset: raise Exception('Model or dataset must be set.') + cookies = None + if self.args.token is not None: + api = HubApi() + _, cookies = api.login( + access_token=self.args.token, save_session=False) if self.args.model: if len(self.args.files) == 1: # download single file model_file_download( @@ -125,7 +136,8 @@ class DownloadCMD(CLICommand): self.args.files[0], cache_dir=self.args.cache_dir, local_dir=self.args.local_dir, - revision=self.args.revision) + revision=self.args.revision, + cookies=cookies) elif len( self.args.files) > 1: # download specified multiple files. snapshot_download( @@ -135,7 +147,7 @@ class DownloadCMD(CLICommand): local_dir=self.args.local_dir, allow_file_pattern=self.args.files, max_workers=self.args.max_workers, - ) + cookies=cookies) else: # download repo snapshot_download( self.args.model, @@ -145,7 +157,7 @@ class DownloadCMD(CLICommand): allow_file_pattern=convert_patterns(self.args.include), ignore_file_pattern=convert_patterns(self.args.exclude), max_workers=self.args.max_workers, - ) + cookies=cookies) elif self.args.dataset: dataset_revision: str = self.args.revision if self.args.revision else DEFAULT_DATASET_REVISION if len(self.args.files) == 1: # download single file @@ -154,7 +166,8 @@ class DownloadCMD(CLICommand): self.args.files[0], cache_dir=self.args.cache_dir, local_dir=self.args.local_dir, - revision=dataset_revision) + revision=dataset_revision, + cookies=cookies) elif len( self.args.files) > 1: # download specified multiple files. dataset_snapshot_download( @@ -164,7 +177,7 @@ class DownloadCMD(CLICommand): local_dir=self.args.local_dir, allow_file_pattern=self.args.files, max_workers=self.args.max_workers, - ) + cookies=cookies) else: # download repo dataset_snapshot_download( self.args.dataset, @@ -174,6 +187,6 @@ class DownloadCMD(CLICommand): allow_file_pattern=convert_patterns(self.args.include), ignore_file_pattern=convert_patterns(self.args.exclude), max_workers=self.args.max_workers, - ) + cookies=cookies) else: pass # noop diff --git a/modelscope/cli/upload.py b/modelscope/cli/upload.py index 453a6314..3aa18908 100644 --- a/modelscope/cli/upload.py +++ b/modelscope/cli/upload.py @@ -91,7 +91,7 @@ class UploadCMD(CLICommand): '--endpoint', type=str, default=get_endpoint(), - help='Endpoint for Modelscope service.') + help='Endpoint for ModelScope service.') parser.set_defaults(func=subparser_func) @@ -137,14 +137,16 @@ class UploadCMD(CLICommand): # Check token and login # The cookies will be reused if the user has logged in before. - api = HubApi(endpoint=self.args.endpoint) - + cookies = None if self.args.token: - api.login(access_token=self.args.token) - cookies = ModelScopeConfig.get_cookies() + api = HubApi(endpoint=self.args.endpoint) + _, cookies = api.login( + access_token=self.args.token, save_session=False) + else: + cookies = ModelScopeConfig.get_cookies() if cookies is None: raise ValueError( - 'The `token` is not provided! ' + 'No credential found for entity upload. ' 'You can pass the `--token` argument, ' 'or use api.login(access_token=`your_sdk_token`). ' 'Your token is available at https://modelscope.cn/my/myaccesstoken' @@ -158,6 +160,7 @@ class UploadCMD(CLICommand): repo_type=self.args.repo_type, commit_message=self.args.commit_message, commit_description=self.args.commit_description, + cookies=cookies, ) elif os.path.isdir(self.local_path): api.upload_folder( @@ -170,6 +173,7 @@ class UploadCMD(CLICommand): allow_patterns=convert_patterns(self.args.include), ignore_patterns=convert_patterns(self.args.exclude), max_workers=self.args.max_workers, + cookies=cookies, ) else: raise ValueError(f'{self.local_path} is not a valid local path') diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index f5a2f39b..89b39546 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -115,6 +115,7 @@ class HubApi: def login( self, access_token: Optional[str] = None, + save_session: Optional[bool] = True ): """Login with your SDK access token, which can be obtained from https://www.modelscope.cn user center. @@ -146,6 +147,9 @@ class HubApi: token = d[API_RESPONSE_FIELD_DATA][API_RESPONSE_FIELD_GIT_ACCESS_TOKEN] cookies = r.cookies + if not save_session: + return None, cookies + # save token and cookie ModelScopeConfig.save_token(token) ModelScopeConfig.save_cookies(cookies)