diff --git a/modelscope/cli/download.py b/modelscope/cli/download.py index fe2b37e9..52786e37 100644 --- a/modelscope/cli/download.py +++ b/modelscope/cli/download.py @@ -26,7 +26,7 @@ class DownloadCMD(CLICommand): """ define args for download command. """ parser: ArgumentParser = parsers.add_parser(DownloadCMD.name) - group = parser.add_mutually_exclusive_group(required=True) + group = parser.add_mutually_exclusive_group() group.add_argument( '--model', type=str, @@ -37,6 +37,20 @@ class DownloadCMD(CLICommand): type=str, help='The id of the dataset to be downloaded. For download, ' 'the id of either a model or dataset must be provided.') + parser.add_argument( + 'repo_id', + type=str, + nargs='?', + default=None, + help='Optional, ' + 'ID of the repo to download, It can also be set by --model or --dataset.' + ) + parser.add_argument( + '--repo-type', + choices=['model', 'dataset'], + default='model', + help="Type of repo to download from (defaults to 'model').", + ) parser.add_argument( '--revision', type=str, @@ -77,6 +91,16 @@ class DownloadCMD(CLICommand): parser.set_defaults(func=subparser_func) def execute(self): + if self.args.repo_id is not None: + if self.args.repo_type == 'model': + self.args.model = self.args.repo_id + elif self.args.repo_type == 'dataset': + self.args.dataset = self.args.repo_id + else: + raise Exception('Not support repo-type: %s' + % self.args.repo_type) + if not self.args.model and not self.args.dataset: + raise Exception('Model or dataset must be set.') if self.args.model: if len(self.args.files) == 1: # download single file model_file_download( diff --git a/modelscope/hub/file_download.py b/modelscope/hub/file_download.py index b5a7a547..8add72bb 100644 --- a/modelscope/hub/file_download.py +++ b/modelscope/hub/file_download.py @@ -168,8 +168,8 @@ def _repo_file_download( if not repo_type: repo_type = REPO_TYPE_MODEL if repo_type not in REPO_TYPE_SUPPORT: - raise InvalidParameter('Invalid repo type: %s, only support: %s' ( - repo_type, REPO_TYPE_SUPPORT)) + raise InvalidParameter('Invalid repo type: %s, only support: %s' % + (repo_type, REPO_TYPE_SUPPORT)) temporary_cache_dir, cache = create_temporary_directory_and_cache( repo_id, local_dir=local_dir, cache_dir=cache_dir, repo_type=repo_type) @@ -471,8 +471,9 @@ def http_get_model_file( with open(temp_file_path, 'rb') as f: partial_length = f.seek(0, io.SEEK_END) progress.update(partial_length) - if partial_length > file_size: + if partial_length >= file_size: break + # closed range[], from 0. get_headers['Range'] = 'bytes=%s-%s' % (partial_length, file_size - 1) with open(temp_file_path, 'ab+') as f: diff --git a/tests/cli/test_download_cmd.py b/tests/cli/test_download_cmd.py index 855ce9ff..4a189833 100644 --- a/tests/cli/test_download_cmd.py +++ b/tests/cli/test_download_cmd.py @@ -57,6 +57,11 @@ class DownloadCMDTest(unittest.TestCase): stat, output = subprocess.getstatusoutput(cmd) self.assertEqual(stat, 0) + def test_download_with_position_arg(self): + cmd = f'python -m modelscope.cli.cli download {self.model_id}' + stat, output = subprocess.getstatusoutput(cmd) + self.assertEqual(stat, 0) + def test_download_with_cache(self): cmd = f'python -m modelscope.cli.cli download --model {self.model_id} --cache_dir {self.tmp_dir}' stat, output = subprocess.getstatusoutput(cmd)