fix

404875ca · Timothy J. Baek · 6dbd59c7 · 404875ca
Commit 404875ca authored May 24, 2024 by Timothy J. Baek
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

backend/main.py backend/main.py +4 -1

No files found.
--- a/backend/main.py
+++ b/backend/main.py
@@ -307,7 +307,10 @@ async def get_models(user=Depends(get_verified_user)):
    for custom_model in custom_models:
        if custom_model.base_model_id == None:
            for model in models:
-                if custom_model.id == model["id"]:
+                if (
+                    custom_model.id == model["id"]
+                    or custom_model.id == model["id"].split(":")[0]
+                ):
                    model["name"] = custom_model.name
                    model["info"] = custom_model.model_dump()
        else: