Skip to content

Commit

Permalink
Fix model download code to better handle when uniqueID and huggingfac…
Browse files Browse the repository at this point in the history
…e_repo are not the same
  • Loading branch information
dadmobile committed Feb 22, 2024
1 parent 850d4bd commit 448643b
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions transformerlab/routers/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def get_model_details_from_gallery(model_id: str):
result = None

for model in gallery:
if model['huggingface_repo'] == model_id or model['uniqueID'] == model_id:
if model['uniqueID'] == model_id or model['huggingface_repo'] == model_id:
result = model
break

Expand Down Expand Up @@ -274,7 +274,9 @@ async def download_model_from_gallery(gallery_id: str, job_id: int | None = None
if gallery_entry is None:
return {"status": "error", "message": "Model not found in gallery"}

return await download_huggingface_model(gallery_id, gallery_entry, job_id)
# Need to use huggingface repo to download - not always the same as uniqueID
huggingface_id = gallery_entry.get('huggingface_repo', gallery_id)
return await download_huggingface_model(huggingface_id, gallery_entry, job_id)


@router.get("/model/get_conversation_template")
Expand Down

0 comments on commit 448643b

Please sign in to comment.