153 lines
4.9 KiB
Python
153 lines
4.9 KiB
Python
import os
|
|
import re
|
|
import torch
|
|
import hashlib
|
|
import urllib.request
|
|
import urllib.error
|
|
from tqdm import tqdm
|
|
from urllib.parse import urlparse
|
|
from typing import Dict, Optional
|
|
|
|
|
|
def natural_sort_key(s, regex=re.compile("([0-9]+)")):
|
|
return [int(text) if text.isdigit() else text.lower() for text in regex.split(s)]
|
|
|
|
|
|
def walk_files(path, allowed_extensions=None):
|
|
if not os.path.exists(path):
|
|
return
|
|
|
|
if allowed_extensions is not None:
|
|
allowed_extensions = set(allowed_extensions)
|
|
|
|
items = list(os.walk(path, followlinks=True))
|
|
items = sorted(items, key=lambda x: natural_sort_key(x[0]))
|
|
|
|
for root, _, files in items:
|
|
for filename in sorted(files, key=natural_sort_key):
|
|
if allowed_extensions is not None:
|
|
_, ext = os.path.splitext(filename)
|
|
if ext not in allowed_extensions:
|
|
continue
|
|
|
|
# Skip hidden files
|
|
if "/." in root or "\\." in root:
|
|
continue
|
|
|
|
yield os.path.join(root, filename)
|
|
|
|
|
|
def load_file_from_url(
|
|
url: str,
|
|
*,
|
|
model_dir: str,
|
|
progress: bool = True,
|
|
file_name: Optional[str] = None,
|
|
) -> str:
|
|
"""Download a file from `url` into `model_dir`, using the file present if possible.
|
|
|
|
Returns the path to the downloaded file.
|
|
"""
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
if not file_name:
|
|
parts = urlparse(url)
|
|
file_name = os.path.basename(parts.path)
|
|
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
|
|
if not os.path.exists(cached_file):
|
|
print(f'Downloading: "{url}" to {cached_file}\n')
|
|
from torch.hub import download_url_to_file
|
|
|
|
download_url_to_file(url, cached_file, progress=progress)
|
|
return cached_file
|
|
|
|
|
|
def calculate_sha(file: str, force=False) -> Optional[str]:
|
|
sha_file = f"{file}.sha"
|
|
|
|
# Check if the .sha file exists
|
|
if not force and os.path.exists(sha_file):
|
|
try:
|
|
with open(sha_file, "r") as f:
|
|
stored_hash = f.read().strip()
|
|
if stored_hash:
|
|
return stored_hash
|
|
except IOError as e:
|
|
print(f"Failed to read hash: {e}")
|
|
|
|
# Calculate the hash if the .sha file doesn't exist or is empty
|
|
try:
|
|
with open(file, "rb") as fp:
|
|
file_hash = hashlib.sha256()
|
|
while chunk := fp.read(8192):
|
|
file_hash.update(chunk)
|
|
calculated_hash = file_hash.hexdigest()
|
|
|
|
# Write the calculated hash to the .sha file
|
|
try:
|
|
with open(sha_file, "w") as f:
|
|
f.write(calculated_hash)
|
|
except IOError as e:
|
|
print(f"Failed to write hash to {sha_file}: {e}")
|
|
|
|
return calculated_hash
|
|
except IOError as e:
|
|
print(f"Failed to read file {file}: {e}")
|
|
return None
|
|
|
|
|
|
def download_file(url: str, dst: str, sha256sum: Optional[str] = None) -> Dict[str, Optional[str]]:
|
|
"""
|
|
Downloads a file from a URL to a destination path, optionally verifying its SHA-256 checksum.
|
|
|
|
:param url: URL of the file to download
|
|
:param dst: Destination path to save the downloaded file
|
|
:param sha256sum: Optional SHA-256 checksum to verify the downloaded file
|
|
:return: Dictionary with file path, download status, calculated checksum, and checksum match status
|
|
"""
|
|
# Ensure the directory exists
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
|
|
file_exists = os.path.isfile(dst)
|
|
file_checksum = None
|
|
checksum_match = None
|
|
downloaded = False
|
|
|
|
try:
|
|
if file_exists:
|
|
file_checksum = calculate_sha(dst)
|
|
if sha256sum:
|
|
checksum_match = file_checksum == sha256sum
|
|
if not checksum_match:
|
|
os.remove(dst)
|
|
|
|
if not file_exists or checksum_match == False:
|
|
with tqdm(unit="B", unit_scale=True, unit_divisor=1024, miniters=1, desc=dst.split("/")[-1]) as t:
|
|
|
|
def reporthook(blocknum, blocksize, totalsize):
|
|
if t.total is None and totalsize > 0:
|
|
t.total = totalsize
|
|
read_so_far = blocknum * blocksize
|
|
t.update(max(0, read_so_far - t.n))
|
|
|
|
urllib.request.urlretrieve(url, dst, reporthook=reporthook)
|
|
downloaded = True
|
|
|
|
file_checksum = calculate_sha(dst, force=True)
|
|
if sha256sum:
|
|
checksum_match = file_checksum == sha256sum
|
|
|
|
except urllib.error.URLError as ex:
|
|
print("Download failed:", ex)
|
|
if os.path.isfile(dst):
|
|
os.remove(dst)
|
|
except Exception as ex:
|
|
print("An error occurred:", ex)
|
|
finally:
|
|
return {"file": dst, "downloaded": downloaded, "sha": file_checksum, "match": checksum_match}
|
|
|
|
|
|
def load_jit_torch_file(model_path: str):
|
|
model = torch.jit.load(model_path, map_location="cpu")
|
|
model.eval()
|
|
return model
|