fix: image model ranking crashes on a non-string search filter (#1898)
This commit is contained in:
@@ -288,7 +288,7 @@ def rank_image_models(system, search=None, sort="fit"):
|
||||
|
||||
for model in IMAGE_MODEL_REGISTRY:
|
||||
# Filter by search
|
||||
if search:
|
||||
if isinstance(search, str) and search:
|
||||
s = search.lower()
|
||||
if s not in model["name"].lower() and s not in model["id"].lower() and s not in model.get("description", "").lower():
|
||||
continue
|
||||
|
||||
15
tests/test_image_models_nonstring_search.py
Normal file
15
tests/test_image_models_nonstring_search.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from services.hwfit.image_models import rank_image_models, IMAGE_MODEL_REGISTRY
|
||||
|
||||
SYS = {"gpu_vram_gb": 0, "has_gpu": False}
|
||||
|
||||
|
||||
def test_rank_image_models_handles_non_string_search():
|
||||
# search is a CLI/API filter arg; a non-string made search.lower() raise
|
||||
# AttributeError. A non-string search should behave as "no filter".
|
||||
out = rank_image_models(SYS, search=123)
|
||||
assert len(out) == len(IMAGE_MODEL_REGISTRY)
|
||||
|
||||
|
||||
def test_rank_image_models_string_filter_still_applies():
|
||||
out = rank_image_models(SYS, search="zzzznotarealmodelzzz")
|
||||
assert out == []
|
||||
Reference in New Issue
Block a user