update token behavior

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

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')

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)