Skip to content

Commit 8a4c629

Browse files
brandonrisinghipsterusername
authored andcommitted
fix: Avoid downloading unsafe .bin files if a safetensors file is available
1 parent a01d44f commit 8a4c629

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

invokeai/backend/model_manager/util/select_hf_files.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
8585
"""Select the proper variant files from a list of HuggingFace repo_id paths."""
8686
result: set[Path] = set()
8787
subfolder_weights: dict[Path, list[SubfolderCandidate]] = {}
88+
safetensors_detected = False
8889
for path in files:
8990
if path.suffix in [".onnx", ".pb", ".onnx_data"]:
9091
if variant == ModelRepoVariant.ONNX:
@@ -119,10 +120,16 @@ def _filter_by_variant(files: List[Path], variant: ModelRepoVariant) -> Set[Path
119120
# We prefer safetensors over other file formats and an exact variant match. We'll score each file based on
120121
# variant and format and select the best one.
121122

123+
if safetensors_detected and path.suffix == ".bin":
124+
continue
125+
122126
parent = path.parent
123127
score = 0
124128

125129
if path.suffix == ".safetensors":
130+
safetensors_detected = True
131+
if parent in subfolder_weights:
132+
subfolder_weights[parent] = [sfc for sfc in subfolder_weights[parent] if sfc.path.suffix != ".bin"]
126133
score += 1
127134

128135
candidate_variant_label = path.suffixes[0] if len(path.suffixes) == 2 else None

0 commit comments

Comments
 (0)