fix pull model name

This commit is contained in:
Michael Yang 2023-06-30 11:54:46 -07:00
parent e69a08f1b1
commit 3ce3caaf65
2 changed files with 10 additions and 10 deletions

View file

@ -151,7 +151,7 @@ def pull(*args, **kwargs):
def run(*args, **kwargs): def run(*args, **kwargs):
name = model.pull(*args, **kwargs) name = model.pull(model_name=kwargs.pop('model'), *args, **kwargs)
kwargs.update({"model": name}) kwargs.update({"model": name})
print(f"Running {name}...") print(f"Running {name}...")
generate(*args, **kwargs) generate(*args, **kwargs)

View file

@ -110,25 +110,25 @@ def download_file(download_url, file_name, file_size):
return local_filename return local_filename
def pull(model, *args, **kwargs): def pull(model_name, *args, **kwargs):
if path.exists(model): if path.exists(model_name):
# a file on the filesystem is being specified # a file on the filesystem is being specified
return model return model_name
# check the remote model location and see if it needs to be downloaded # check the remote model location and see if it needs to be downloaded
url = model url = model_name
file_name = "" file_name = ""
if not validators.url(url) and not url.startswith('huggingface.co'): if not validators.url(url) and not url.startswith('huggingface.co'):
url = get_url_from_directory(model) url = get_url_from_directory(model_name)
file_name = model file_name = model_name
if not (url.startswith('http://') or url.startswith('https://')): if not (url.startswith('http://') or url.startswith('https://')):
url = f'https://{url}' url = f'https://{url}'
if not validators.url(url): if not validators.url(url):
if model in models(MODELS_CACHE_PATH): if model_name in models(MODELS_CACHE_PATH):
# the model is already downloaded, and specified by name # the model is already downloaded, and specified by name
return model return model_name
raise Exception(f'Unknown model {model}') raise Exception(f'Unknown model {model_name}')
local_filename = download_from_repo(url, file_name) local_filename = download_from_repo(url, file_name)