diff --git a/modelscope/cli/create.py b/modelscope/cli/create.py index 3285ab91..76e67e43 100644 --- a/modelscope/cli/create.py +++ b/modelscope/cli/create.py @@ -3,8 +3,8 @@ from argparse import ArgumentParser, _SubParsersAction from modelscope.cli.base import CLICommand from modelscope.hub.api import HubApi -from modelscope.hub.constants import (Licenses, ModelVisibility, Visibility, - VisibilityMap) +from modelscope.hub.constants import (GatedMode, Licenses, ModelVisibility, + Visibility, VisibilityMap) from modelscope.hub.utils.aigc import AigcModel from modelscope.hub.utils.utils import resolve_endpoint from modelscope.utils.constant import REPO_TYPE_MODEL, REPO_TYPE_SUPPORT @@ -82,6 +82,20 @@ class CreateCMD(CLICommand): help= 'If True, do not raise error when repo already exists. Defaults to False.', ) + parser.add_argument( + '--gated', + dest='gated_mode', + action='store_true', + default=None, + help= + 'Enable gated mode (application-based download) for private repos.', + ) + parser.add_argument( + '--no-gated', + dest='gated_mode', + action='store_false', + help='Disable gated mode for private repos (normal private).', + ) parser.add_argument( '--endpoint', type=str, @@ -176,6 +190,7 @@ class CreateCMD(CLICommand): exist_ok=self.args.exist_ok, create_default_config=True, endpoint=endpoint, + gated_mode=self.args.gated_mode, ) def _create_aigc_model(self): @@ -225,7 +240,8 @@ class CreateCMD(CLICommand): visibility=visibility_idx, license=self.args.license, chinese_name=self.args.chinese_name, - aigc_model=aigc_model) + aigc_model=aigc_model, + gated_mode=self.args.gated_mode) print(f'Successfully created AIGC model: {model_url}') except Exception as e: print(f'Error creating AIGC model: {e}') diff --git a/modelscope/hub/api.py b/modelscope/hub/api.py index 8a7b8409..f223eed0 100644 --- a/modelscope/hub/api.py +++ b/modelscope/hub/api.py @@ -329,7 +329,8 @@ class HubApi: original_model_id: Optional[str] = '', endpoint: Optional[str] = None, token: Optional[str] = None, - aigc_model: Optional['AigcModel'] = None) -> str: + aigc_model: Optional['AigcModel'] = None, + gated_mode: Optional[bool] = None) -> str: """Create model repo at ModelScope Hub. Args: @@ -343,6 +344,9 @@ class HubApi: aigc_model (AigcModel, optional): AigcModel instance for AIGC model creation. If provided, will create an AIGC model with automatic file upload. Refer to modelscope.hub.utils.aigc.AigcModel for details. + gated_mode (bool, optional): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is PRIVATE (1). Returns: str: URL of the created model repository @@ -374,6 +378,12 @@ class HubApi: 'TrainId': os.environ.get('MODELSCOPE_TRAIN_ID', '') } + if gated_mode is not None: + if visibility != ModelVisibility.PRIVATE: + logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') + else: + body['ProtectedMode'] = 1 if gated_mode else 2 + # Set path based on model type if aigc_model is not None: # Use AIGC model endpoint @@ -1428,7 +1438,8 @@ class HubApi: visibility: Optional[int] = DatasetVisibility.PUBLIC, description: Optional[str] = '', endpoint: Optional[str] = None, - token: Optional[str] = None) -> str: + token: Optional[str] = None, + gated_mode: Optional[bool] = None) -> str: """ Create a dataset in ModelScope. @@ -1441,6 +1452,9 @@ class HubApi: description (str, optional): The description of the dataset. Defaults to ''. endpoint (str, optional): The endpoint to use. If not provided, the default endpoint is used. token (str, optional): The access token for authentication. + gated_mode (bool, optional): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is PRIVATE (1). Returns: str: The URL of the created dataset repository. @@ -1462,6 +1476,12 @@ class HubApi: 'Description': (None, description) } + if gated_mode is not None: + if visibility != DatasetVisibility.PRIVATE: + logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.') + else: + files['ProtectedMode'] = (None, 1 if gated_mode else 2) + r = self.session.post( path, files=files, @@ -2200,6 +2220,7 @@ class HubApi: exist_ok: Optional[bool] = False, create_default_config: Optional[bool] = True, aigc_model: Optional[AigcModel] = None, + gated_mode: Optional[bool] = None, **kwargs, ) -> str: """ @@ -2217,6 +2238,9 @@ class HubApi: In the format of `https://www.modelscope.cn` or 'https://www.modelscope.ai' exist_ok (Optional[bool]): If the repo exists, whether to return the repo url directly. create_default_config (Optional[bool]): If True, create a default configuration file in the model repo. + gated_mode (Optional[bool]): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is ``private``. **kwargs: The additional arguments. Returns: @@ -2256,6 +2280,7 @@ class HubApi: aigc_model=aigc_model, token=token, endpoint=endpoint, + gated_mode=gated_mode, ) if create_default_config: with tempfile.TemporaryDirectory() as temp_cache_dir: @@ -2290,6 +2315,7 @@ class HubApi: visibility=visibility, token=token, endpoint=endpoint, + gated_mode=gated_mode, ) print(f'New dataset created successfully at {repo_url}.', flush=True) @@ -3882,7 +3908,8 @@ class HubApi: repo_id: str, repo_type: Literal['model', 'dataset'], visibility: Literal['private', 'public'], - token: Union[str, None] = None + token: Union[str, None] = None, + gated_mode: Optional[bool] = None, ) -> dict: """ Set the visibility of a repo. @@ -3893,6 +3920,9 @@ class HubApi: visibility (Literal['private', 'public']): The visibility to set, `private` or `public`. token (Union[str, None]): The access token. If None, will use the cookies from the local cache. See `https://modelscope.cn/my/myaccesstoken` to get your token. + gated_mode (Optional[bool]): Gated mode for private repos. + True = gated (application-based download), False = off (normal private). + Only effective when visibility is ``private``. Returns: dict: The response from the server. @@ -3907,6 +3937,10 @@ class HubApi: visibility_code: int = visibility_map.get(visibility, 5) cookies = self.get_cookies(access_token=token, cookies_required=True) + if gated_mode is not None and visibility != 'private': + logger.warning('gated_mode is only effective when visibility is private, ignored.') + gated_mode = None + if repo_type == REPO_TYPE_MODEL: model_info = self.get_model(model_id=repo_id, token=token) path = f'{self.endpoint}/api/v1/models/{repo_id}' @@ -3916,11 +3950,15 @@ class HubApi: first = tasks[0] if isinstance(first, dict) and first: model_tasks = first.get('name') + if gated_mode is not None: + pm = 1 if gated_mode else 2 + else: + pm = model_info.get('ProtectedMode', 2) payload = { 'ChineseName': model_info.get('ChineseName', ''), 'ModelFramework': model_info.get('ModelFramework', 'Pytorch'), 'Visibility': visibility_code, - 'ProtectedMode': 2, + 'ProtectedMode': pm, 'ApprovalMode': model_info.get('ApprovalMode', 2), 'Description': model_info.get('Description', ''), 'AigcType': model_info.get('AigcType', ''), @@ -3947,7 +3985,7 @@ class HubApi: path = f'{self.endpoint}/api/v1/datasets/{dataset_idx}' payload = { 'Visibility': visibility_code, - 'ProtectedMode': 2, + 'ProtectedMode': (1 if gated_mode else 2) if gated_mode is not None else 2, } else: raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}') diff --git a/modelscope/hub/constants.py b/modelscope/hub/constants.py index 85ad69d4..c3ced124 100644 --- a/modelscope/hub/constants.py +++ b/modelscope/hub/constants.py @@ -169,6 +169,18 @@ class Visibility(object): PUBLIC = 'public' +class GatedMode(object): + """Gated mode for private repositories. + + Only effective when Visibility is PRIVATE. + API payload key: ``ProtectedMode``. + Values: True = gated (application-based download), + False = off (normal private). + """ + GATED = True + OFF = False + + VisibilityMap = { ModelVisibility.PRIVATE: Visibility.PRIVATE, ModelVisibility.INTERNAL: Visibility.INTERNAL,