Spaces:
Runtime error
Runtime error
File size: 2,787 Bytes
578fa37 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
import json
import requests
# fetch_pretrains returns an array with the names
# download_pretrain downloads the pretrain with the name
PRETRAINS_URL = "https://gist.githubusercontent.com/miqu-s/35042d5eaaebe311ecc2530040434192/raw/54fe5b54723bae3182479bc60d847b9dcb17270a/pretrains"
def fetch_pretrains():
try:
response = requests.get(PRETRAINS_URL)
if response.status_code == 200:
pretrains_data = json.loads(response.text)
return list(pretrains_data.keys())
else:
print("Failed to fetch pretrains data. Status code:", response.status_code)
return []
except Exception as e:
print("An error occurred while fetching pretrains data:", e)
return []
def download_pretrain(pretrain_name):
try:
response = requests.get(PRETRAINS_URL)
if response.status_code == 200:
pretrains_data = json.loads(response.text)
if pretrain_name in pretrains_data:
pretrain_urls = pretrains_data[pretrain_name]
total_files = len(pretrain_urls)
print(f"Downloading {total_files} files for {pretrain_name}:")
for index, (key, url) in enumerate(pretrain_urls.items(), 1):
filename = f"{pretrain_name}_{key}.pth"
download_file(url, filename, index, total_files)
else:
print("Pretrain not found.")
else:
print("Failed to fetch pretrains data. Status code:", response.status_code)
except Exception as e:
print("An error occurred while downloading pretrain:", e)
def download_file(url, filename, current_file, total_files):
try:
response = requests.get(url, stream=True)
if response.status_code == 200:
directory = "assets/pretrains_v2"
if not os.path.exists(directory):
os.makedirs(directory)
filepath = os.path.join(directory, filename)
with open(filepath, 'wb') as f:
total_length = int(response.headers.get('content-length'))
print(f"Downloading {filename} [{current_file}/{total_files}]")
dl = 0
for data in response.iter_content(chunk_size=4096):
dl += len(data)
f.write(data)
done = int(50 * dl / total_length)
print("\r[%s%s]" % ('=' * done, ' ' * (50 - done)), end='', flush=True)
print(f"\rDownloaded {filename} successfully.")
else:
print(f"Failed to download {filename}. Status code:", response.status_code)
except Exception as e:
print(f"An error occurred while downloading {filename}:", e) |