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:
|
for model in IMAGE_MODEL_REGISTRY:
|
||||||
# Filter by search
|
# Filter by search
|
||||||
if search:
|
if isinstance(search, str) and search:
|
||||||
s = search.lower()
|
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():
|
if s not in model["name"].lower() and s not in model["id"].lower() and s not in model.get("description", "").lower():
|
||||||
continue
|
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