HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix NCNN by using safe_download() (#11761)

Browse source
This commit is contained in:
Glenn Jocher 2024-05-08 15:23:22 +02:00 • committed by GitHub
parent 4527690a7d
commit 4350babc75
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

16
ultralytics/engine/exporter.py
View file

@ -533,7 +533,7 @@ class Exporter:
f_ts = self.file.with_suffix(".torchscript") f_ts = self.file.with_suffix(".torchscript")
name = Path("pnnx.exe" if WINDOWS else "pnnx") # PNNX filename name = Path("pnnx.exe" if WINDOWS else "pnnx") # PNNX filename
pnnx = name if name.is_file() else ROOT / name pnnx = name if name.is_file() else (ROOT / name)
if not pnnx.is_file(): if not pnnx.is_file():
LOGGER.warning( LOGGER.warning(
f"{prefix} WARNING ⚠️ PNNX not found. Attempting to download binary file from " f"{prefix} WARNING ⚠️ PNNX not found. Attempting to download binary file from "
@ -542,21 +542,19 @@ class Exporter:
) )
system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux" system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux"
try: try:
_, assets = get_github_assets(repo="pnnx/pnnx") release, assets = get_github_assets(repo="pnnx/pnnx")
asset = [x for x in assets if f"{system}.zip" in x][0] asset = [x for x in assets if f"{system}.zip" in x][0]
assert isinstance(asset, str), "Unable to retrieve PNNX repo assets" # i.e. pnnx-20240410-macos.zip assert isinstance(asset, str), "Unable to retrieve PNNX repo assets" # i.e. pnnx-20240410-macos.zip
LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}") LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}")
asset = attempt_download_asset(asset, repo="pnnx/pnnx", release="latest")
except Exception as e: except Exception as e:
url = f"https://github.com/pnnx/pnnx/releases/download/20240410/pnnx-20240410-{system}.zip" release = "20240410"
LOGGER.warning(f"{prefix} WARNING ⚠️ PNNX GitHub assets not found: {e}, using default {url}") asset = f"pnnx-{release}-{system}.zip"
asset = safe_download(url) LOGGER.warning(f"{prefix} WARNING ⚠️ PNNX GitHub assets not found: {e}, using default {asset}")
if check_is_path_safe(Path.cwd(), asset): # avoid path traversal security vulnerability unzip_dir = safe_download(f"https://github.com/pnnx/pnnx/releases/download/{release}/{asset}", delete=True)
unzip_dir = Path(asset).with_suffix("") if check_is_path_safe(Path.cwd(), unzip_dir): # avoid path traversal security vulnerability
(unzip_dir / name).rename(pnnx) # move binary to ROOT (unzip_dir / name).rename(pnnx) # move binary to ROOT
pnnx.chmod(0o777) # set read, write, and execute permissions for everyone pnnx.chmod(0o777) # set read, write, and execute permissions for everyone
shutil.rmtree(unzip_dir) # delete unzip dir shutil.rmtree(unzip_dir) # delete unzip dir
Path(asset).unlink(missing_ok=True) # delete zip
ncnn_args = [ ncnn_args = [
f'ncnnparam={f / "model.ncnn.param"}', f'ncnnparam={f / "model.ncnn.param"}',