diff --git a/simple_syrup/runtime/ultralytics_loader.py b/simple_syrup/runtime/ultralytics_loader.py index fee0f11..183ea3a 100644 --- a/simple_syrup/runtime/ultralytics_loader.py +++ b/simple_syrup/runtime/ultralytics_loader.py @@ -93,18 +93,39 @@ class UltralyticsLoaderService: ) def model_choices(self) -> list[str]: - """Return curated and local Ultralytics model choices for ComfyUI dropdowns.""" + """Return installed choices first, followed by downloadable catalog choices.""" self._register_model_folders() curated_choices = self._choice_service.ultralytics_choices() - curated_local_paths = { - _catalog_selection(entry) for entry in ULTRALYTICS_ENTRIES + catalog_choice_labels = { + _catalog_selection(entry): entry.display_name + for entry in ULTRALYTICS_ENTRIES } - choices = curated_choices + [ - choice - for choice in self.available_models() - if choice not in curated_local_paths + available_choices = self.available_models() + visible_catalog_choices = set(curated_choices) + installed_catalog_choices = [ + entry.display_name + for entry in ULTRALYTICS_ENTRIES + if ( + entry.display_name in visible_catalog_choices + and _catalog_selection(entry) in available_choices + ) ] + installed_non_catalog_choices = [ + choice + for choice in available_choices + if choice not in catalog_choice_labels + ] + downloadable_choices = [ + choice + for choice in curated_choices + if choice not in installed_catalog_choices + ] + choices = ( + installed_non_catalog_choices + + installed_catalog_choices + + downloadable_choices + ) return choices or [NO_LOCAL_ULTRALYTICS_MODELS] def available_models(self) -> list[str]: diff --git a/tests/test_ultralytics_loader.py b/tests/test_ultralytics_loader.py index a462f5c..f2c6f38 100644 --- a/tests/test_ultralytics_loader.py +++ b/tests/test_ultralytics_loader.py @@ -85,6 +85,32 @@ def test_model_choices_include_curated_downloadable_models(tmp_path: Path) -> No assert not any("Score" in choice for choice in choices) +def test_model_choices_list_installed_models_before_downloadable_entries( + tmp_path: Path, +) -> None: + """Installed choices precede curated models that still require a download.""" + + models_dir = tmp_path / "models" + bbox_dir = models_dir / "ultralytics" / "bbox" + bbox_dir.mkdir(parents=True) + (bbox_dir / "face_yolov8n_v2.pt").write_bytes(b"checkpoint") + (bbox_dir / "local-detector.pt").write_bytes(b"checkpoint") + folder_paths = _folder_paths(models_dir) + service = UltralyticsLoaderService( + folder_paths_module=folder_paths, + choice_service=_choice_service(show_downloadable_models=True), + ) + + choices = service.model_choices() + + assert choices[:2] == [ + "bbox/local-detector.pt", + "Bingsu Face YOLOv8n v2 (6.23MB)", + ] + assert choices[2] == "Anzhc Face -seg (6.52MB)" + assert "bbox/face_yolov8n_v2.pt" not in choices + + def test_hidden_catalog_choices_exclude_installed_curated_model( tmp_path: Path, ) -> None: