Coverage for src/bioimageio/spec/_hf.py: 98%

1from __future__ import annotations 

2 

3import os 

4import tempfile 

5import warnings 

6from contextlib import nullcontext 

7from functools import cache 

8from pathlib import Path 

9 

10from loguru import logger 

11 

12from bioimageio.spec import save_bioimageio_package_as_folder 

13from bioimageio.spec._internal.validation_context import get_validation_context 

14from bioimageio.spec.model.v0_5 import ModelDescr 

15 

16from ._hf_card import create_huggingface_model_card 

17from ._version import VERSION 

18 

19 

20@cache 

21def get_huggingface_api(): # pragma: no cover 

22 from huggingface_hub import HfApi 

23 

24 return HfApi(library_name="bioimageio.spec", library_version=VERSION) 

25 

26 

27def push_to_hub( 

28 descr: ModelDescr, 

29 username_or_org: str, 

30 *, 

31 prep_dir: os.PathLike[str] | str | None = None, 

32 prep_only_no_upload: bool = False, 

33 create_pr: bool | None = None, 

34): 

35 """Push the model package described by `descr` to the Hugging Face Hub. 

36 

37 Note: 

38 - Uses `descr.id` as the repository name under the provided `username_or_org`. 

39 - If `descr.version` is set, the model package is uploaded to the 'main' branch 

40 and tagged with the version. 

41 - If `descr.version` is `None`, the model package is uploaded to the 'draft' branch. 

42 

43 Args: 

44 descr: The model description to be pushed to the Hugging Face Hub. 

45 username_or_org: The Hugging Face username or organization under which the model package will be uploaded. 

46 The model ID from `descr.id` will be used as the repository name. 

47 prep_dir: Optional path to an empty directory where the model package will be prepared before uploading. 

48 prep_only_no_upload: If `True`, only prepare the model package in `prep_dir` without uploading it 

49 to the Hugging Face Hub. 

50 create_pr: If `False` commit directly to the 'main'/'draft' branch. 

51 If `True`, create a pull request targeting 'main'/'draft'. 

52 Defaults to `True` if uploading to a model description with version (to the main branch), 

53 and `False` if uploading a model description without version (to the 'draft' branch). 

54 

55 Examples: 

56 Upload a model description as a new version to the main branch 

57 (`descr.id` and `descr.version` must be set): 

58 

59 >>> descr = ModelDescr(id="my-model-id", version="1.0", create_pr=False, ...) # doctest: +SKIP 

60 >>> push_to_hub(descr, "my_hf_username") # doctest: +SKIP 

61 

62 Upload a model description as a draft to the 'draft' branch 

63 (`descr.id` must be set; `descr.version` must be `None`): 

64 

65 >>> descr = ModelDescr(id="my-model-id", version=None, ...) # doctest: +SKIP 

66 >>> push_to_hub(descr, "my_hf_username") # doctest: +SKIP 

67 

68 See what would be uploaded without actually uploading: 

69 

70 >>> push_to_hub(..., prep_dir="empty_local_folder", prep_only_no_upload=True) # doctest: +SKIP 

71 

72 """ 

73 

74 if descr.id is None: 

75 raise ValueError("descr.id must be set to push to Hugging Face Hub.") 

76 repo_id = f"{username_or_org}/{descr.id}" 

77 

78 if prep_dir is None: 

79 ctxt = tempfile.TemporaryDirectory(suffix="_" + repo_id.replace("/", "_")) 

80 elif Path(prep_dir).exists() and any(Path(prep_dir).iterdir()): 

81 raise ValueError("Provided `prep_dir` is not empty.") 

82 # TODO: implement resuming upload 

83 # prep_dir: If a non-empty folder is provided, it will be attempted to continue an interrupted upload. 

84 # logger.info(f"Continuing upload from {prep_dir}") 

85 # if prep_only_no_upload: 

86 # raise ValueError("`prep_only_no_upload` is True but `prep_dir` is non-empty.") 

87 else: 

88 ctxt = nullcontext(prep_dir) 

89 

90 with ctxt as pdir: 

91 _push_to_hub_impl( 

92 descr, 

93 repo_id=repo_id, 

94 prep_dir=Path(pdir), 

95 prep_only=prep_only_no_upload, 

96 create_pr=create_pr, 

97 ) 

98 

99 

100def _push_to_hub_impl( 

101 descr: ModelDescr, 

102 *, 

103 repo_id: str, 

104 prep_dir: Path, 

105 prep_only: bool, 

106 create_pr: bool | None, 

107): 

108 readme, referenced_files = create_huggingface_model_card(descr, repo_id=repo_id) 

109 referenced_files_subfolders = {"images"} 

110 assert not ( 

111 unexpected := [ 

112 rf 

113 for rf in referenced_files 

114 if not any(rf.startswith(f"{sf}/") for sf in referenced_files_subfolders) 

115 ] 

116 ), f"unexpected folder of referenced files: {unexpected}" 

117 

118 logger.info(f"Preparing model for upload at {prep_dir}.") 

119 prep_dir.mkdir(parents=True, exist_ok=True) 

120 _ = (prep_dir / "README.md").write_text(readme, encoding="utf-8") 

121 for img_name, img_data in referenced_files.items(): 

122 image_path = prep_dir / img_name 

123 image_path.parent.mkdir(parents=True, exist_ok=True) 

124 _ = image_path.write_bytes(img_data) 

125 

126 with get_validation_context().replace(file_name="bioimageio.yaml"): 

127 _ = save_bioimageio_package_as_folder(descr, output_path=prep_dir / "package") 

128 

129 logger.info(f"Prepared model for upload at {prep_dir}") 

130 

131 commit_message = f"Upload {descr.version or 'draft'} with bioimageio.spec {VERSION}" 

132 commit_description = ( 

133 f"Version comment: {descr.version_comment}" if descr.version_comment else None 

134 ) 

135 

136 if not prep_only: # pragma: no cover 

137 logger.info(f"Pushing model '{descr.id}' to Hugging Face Hub") 

138 

139 api = get_huggingface_api() 

140 repo_url = api.create_repo(repo_id=repo_id, exist_ok=True, repo_type="model") 

141 logger.info(f"Created repository at {repo_url}") 

142 

143 existing_refs = api.list_repo_refs( 

144 repo_id=repo_id, repo_type="model", include_pull_requests=True 

145 ) 

146 has_draft_ref = False 

147 has_tag = False 

148 for ref in existing_refs.branches + existing_refs.tags: 

149 if ref.name == str(descr.version): 

150 has_tag = True 

151 if ref.name == "draft": 

152 has_draft_ref = True 

153 

154 if descr.version is None: 

155 revision = "draft" 

156 if not has_draft_ref: 

157 api.create_branch(repo_id=repo_id, branch="draft", repo_type="model") 

158 else: 

159 revision = None 

160 

161 if create_pr is None: 

162 # default to creating a PR if commiting to main branch, 

163 # commit directly to 'draft' branch 

164 create_pr = revision is None 

165 

166 commit_info = api.upload_folder( 

167 repo_id=repo_id, 

168 revision=revision, 

169 folder_path=prep_dir, 

170 delete_patterns=[f"{sf}/*" for sf in referenced_files_subfolders] 

171 + ["package/*"], 

172 commit_message=commit_message, 

173 commit_description=commit_description, 

174 create_pr=create_pr, 

175 ) 

176 logger.info(f"Created commit {commit_info.commit_url}") 

177 if descr.version is not None: 

178 if has_tag: 

179 warnings.warn(f"Moving existing version tag {descr.version}.") 

180 

181 api.create_tag( 

182 repo_id=repo_id, 

183 tag=str(descr.version), 

184 revision=commit_info.oid, 

185 tag_message=descr.version_comment, 

186 exist_ok=True, 

187 )