From 19d247b573b760b4bdaa2291023a70e5733160d6 Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Tue, 27 Dec 2022 18:20:07 +0000 Subject: [PATCH 1/6] Add shutil to setup.py --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 3c58cbf..18ef841 100644 --- a/setup.py +++ b/setup.py @@ -12,6 +12,7 @@ install_requires=[ "filelock", "Pillow", + "shutil", "torch", "fire", "humanize", From 6626f7347d4cf551ffeec71026b3a5f2c656aa59 Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Tue, 27 Dec 2022 18:20:51 +0000 Subject: [PATCH 2/6] Fetch file cached function refactor for efficiency --- point_e/models/download.py | 42 ++++++++++++++++++++++++++------------ 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/point_e/models/download.py b/point_e/models/download.py index b1760e1..e7706cd 100644 --- a/point_e/models/download.py +++ b/point_e/models/download.py @@ -8,6 +8,7 @@ import requests import torch +import shutil from filelock import FileLock from tqdm.auto import tqdm @@ -39,26 +40,41 @@ def fetch_file_cached( """ if cache_dir is None: cache_dir = default_cache_dir() + os.makedirs(cache_dir, exist_ok=True) local_path = os.path.join(cache_dir, url.split("/")[-1]) + + # If local path exists return it if os.path.exists(local_path): return local_path + # Download the file using the requests library response = requests.get(url, stream=True) size = int(response.headers.get("content-length", "0")) - with FileLock(local_path + ".lock"): - if progress: - pbar = tqdm(total=size, unit="iB", unit_scale=True) - tmp_path = local_path + ".tmp" - with open(tmp_path, "wb") as f: - for chunk in response.iter_content(chunk_size): - if progress: - pbar.update(len(chunk)) - f.write(chunk) - os.rename(tmp_path, local_path) - if progress: - pbar.close() - return local_path + + # Use a progress bar if desired + if progress: + pbar = tqdm(total=size, unit="iB", unit_scale=True) + + # Create a temporary file and download the file in chunks + tmp_path = local_path + ".tmp" + with open(tmp_path, "wb") as f: + for chunk in response.iter_content(chunk_size=chunk_size): + if progress: + pbar.update(len(chunk)) + f.write(chunk) + + # Use shutil to copy the file from the temporary location to the final destination + shutil.copyfile(tmp_path, local_path) + + # Clean up the temporary file + os.remove(tmp_path) + + # Hide progress bar + if progress: + pbar.close() + + return local_path def load_checkpoint( From 936aaa88cd0a7e913784d5a7251f9ad951192be2 Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Tue, 27 Dec 2022 18:24:48 +0000 Subject: [PATCH 3/6] Add back file lock --- point_e/models/download.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/point_e/models/download.py b/point_e/models/download.py index e7706cd..d069c77 100644 --- a/point_e/models/download.py +++ b/point_e/models/download.py @@ -52,17 +52,19 @@ def fetch_file_cached( response = requests.get(url, stream=True) size = int(response.headers.get("content-length", "0")) - # Use a progress bar if desired - if progress: - pbar = tqdm(total=size, unit="iB", unit_scale=True) - - # Create a temporary file and download the file in chunks - tmp_path = local_path + ".tmp" - with open(tmp_path, "wb") as f: - for chunk in response.iter_content(chunk_size=chunk_size): - if progress: - pbar.update(len(chunk)) - f.write(chunk) + # Lock before opening / creating file + with FileLock(local_path + ".lock"): + # Use a progress bar if desired + if progress: + pbar = tqdm(total=size, unit="iB", unit_scale=True) + + # Create a temporary file and download the file in chunks + tmp_path = local_path + ".tmp" + with open(tmp_path, "wb") as f: + for chunk in response.iter_content(chunk_size=chunk_size): + if progress: + pbar.update(len(chunk)) + f.write(chunk) # Use shutil to copy the file from the temporary location to the final destination shutil.copyfile(tmp_path, local_path) @@ -74,6 +76,7 @@ def fetch_file_cached( if progress: pbar.close() + # Return local path return local_path From 0b1b4ffabfed8f1806aa873053183595ba4b058d Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Tue, 27 Dec 2022 18:25:56 +0000 Subject: [PATCH 4/6] Remove chunksize= from param list --- point_e/models/download.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/point_e/models/download.py b/point_e/models/download.py index d069c77..e1062eb 100644 --- a/point_e/models/download.py +++ b/point_e/models/download.py @@ -61,7 +61,7 @@ def fetch_file_cached( # Create a temporary file and download the file in chunks tmp_path = local_path + ".tmp" with open(tmp_path, "wb") as f: - for chunk in response.iter_content(chunk_size=chunk_size): + for chunk in response.iter_content(chunk_size): if progress: pbar.update(len(chunk)) f.write(chunk) From 20e9a94cf9eb7369a3b17b3614ec64ce22f70442 Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Wed, 28 Dec 2022 16:03:20 +0000 Subject: [PATCH 5/6] Remove shutil from setup, thanks dancergraham ! --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index 18ef841..3c58cbf 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,6 @@ install_requires=[ "filelock", "Pillow", - "shutil", "torch", "fire", "humanize", From cb43a9edab9f5a0276e6d1aa2c321f241f41c94b Mon Sep 17 00:00:00 2001 From: Robert McMenemy Date: Wed, 28 Dec 2022 16:38:35 +0000 Subject: [PATCH 6/6] Remove unnecessary comments --- point_e/models/download.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/point_e/models/download.py b/point_e/models/download.py index e1062eb..a621735 100644 --- a/point_e/models/download.py +++ b/point_e/models/download.py @@ -44,39 +44,31 @@ def fetch_file_cached( os.makedirs(cache_dir, exist_ok=True) local_path = os.path.join(cache_dir, url.split("/")[-1]) - # If local path exists return it if os.path.exists(local_path): return local_path - # Download the file using the requests library response = requests.get(url, stream=True) size = int(response.headers.get("content-length", "0")) - # Lock before opening / creating file with FileLock(local_path + ".lock"): - # Use a progress bar if desired if progress: pbar = tqdm(total=size, unit="iB", unit_scale=True) - # Create a temporary file and download the file in chunks tmp_path = local_path + ".tmp" + with open(tmp_path, "wb") as f: for chunk in response.iter_content(chunk_size): if progress: pbar.update(len(chunk)) f.write(chunk) - - # Use shutil to copy the file from the temporary location to the final destination + shutil.copyfile(tmp_path, local_path) - - # Clean up the temporary file + os.remove(tmp_path) - - # Hide progress bar + if progress: pbar.close() - - # Return local path + return local_path