Skip to content

Commit

Permalink
Use custom download_url instead of torchvision's function
Browse files Browse the repository at this point in the history
- importing torchvision is quiet timed consuming
  • Loading branch information
ltdrdata committed Aug 1, 2024
1 parent 707ec9d commit d7f5f08
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion git_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,34 @@
import re
import json
import yaml
from torchvision.datasets.utils import download_url
import requests
from tqdm.auto import tqdm
from git.remote import RemoteProgress


def download_url(url, dest_folder, filename=None):
# Ensure the destination folder exists
if not os.path.exists(dest_folder):
os.makedirs(dest_folder)

# Extract filename from URL if not provided
if filename is None:
filename = os.path.basename(url)

# Full path to save the file
dest_path = os.path.join(dest_folder, filename)

# Download the file
response = requests.get(url, stream=True)
if response.status_code == 200:
with open(dest_path, 'wb') as file:
for chunk in response.iter_content(chunk_size=1024):
if chunk:
file.write(chunk)
else:
print(f"Failed to download file from {url}")


config_path = os.path.join(os.path.dirname(__file__), "config.ini")
nodelist_path = os.path.join(os.path.dirname(__file__), "custom-node-list.json")
working_directory = os.getcwd()
Expand Down

0 comments on commit d7f5f08

Please sign in to comment.