Allow downloading model from HF branch via UI (#1662)

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Tom Jobbins 2023-05-05 17:59:01 +01:00 committed by GitHub
parent 849ad04c96
commit 876fbb97c0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -183,9 +183,9 @@ def count_tokens(text):
def download_model_wrapper(repo_id): def download_model_wrapper(repo_id):
try: try:
downloader = importlib.import_module("download-model") downloader = importlib.import_module("download-model")
repo_id_parts = repo_id.split(":")
model = repo_id model = repo_id_parts[0] if len(repo_id_parts) > 0 else repo_id
branch = "main" branch = repo_id_parts[1] if len(repo_id_parts) > 1 else "main"
check = False check = False
yield ("Cleaning up the model/branch names") yield ("Cleaning up the model/branch names")
@ -370,7 +370,7 @@ def create_model_menus():
with gr.Row(): with gr.Row():
with gr.Column(): with gr.Column():
shared.gradio['custom_model_menu'] = gr.Textbox(label="Download custom model or LoRA", info="Enter Hugging Face username/model path, e.g: facebook/galactica-125m") shared.gradio['custom_model_menu'] = gr.Textbox(label="Download custom model or LoRA", info="Enter the Hugging Face username/model path, for instance: facebook/galactica-125m. To specify a branch, add it at the end after a \":\" character like this: facebook/galactica-125m:main")
shared.gradio['download_model_button'] = gr.Button("Download") shared.gradio['download_model_button'] = gr.Button("Download")
with gr.Column(): with gr.Column():