Spaces:
Running
on
Zero
Running
on
Zero
import os | |
from huggingface_hub import hf_hub_download, list_repo_files | |
import tarfile | |
import shutil | |
from tqdm import tqdm | |
def download_and_extract_file(repo_id, filename, local_dir): | |
""" | |
Download and extract a file from Hugging Face. | |
Args: | |
repo_id (str): The repository ID on Hugging Face | |
filename (str): The name of the file to download | |
local_dir (str): The directory to extract files to | |
""" | |
print(f"Starting download of {filename}...") | |
# Create directory if it doesn't exist | |
os.makedirs(local_dir, exist_ok=True) | |
try: | |
# Download the file using hf_hub_download | |
file_path = hf_hub_download( | |
repo_id=repo_id, | |
filename=filename, | |
repo_type="space", | |
local_dir=local_dir | |
) | |
print(f"Download completed. File saved to: {file_path}") | |
# Extract the tar.xz file with progress bar | |
print(f"Extracting {filename}...") | |
with tarfile.open(file_path, 'r:xz') as tar: | |
# Get total number of files | |
total_files = len(tar.getmembers()) | |
# Extract with progress bar | |
for member in tqdm(tar.getmembers(), total=total_files, desc=f"Extracting {filename}"): | |
tar.extract(member, path=local_dir) | |
print(f"Extraction of {filename} completed successfully!") | |
# Clean up the downloaded archive | |
os.remove(file_path) | |
print(f"Cleaned up downloaded archive {filename}.") | |
except Exception as e: | |
print(f"An error occurred while processing {filename}: {str(e)}") | |
# Clean up in case of error | |
if os.path.exists(file_path): | |
os.remove(file_path) | |
raise | |
def download_video_files(repo_id, local_dir): | |
""" | |
Download all video files from the downloads directory. | |
Args: | |
repo_id (str): The repository ID on Hugging Face | |
local_dir (str): The directory to save videos to | |
""" | |
print("Starting download of video files...") | |
# Create downloads directory if it doesn't exist | |
os.makedirs(local_dir, exist_ok=True) | |
try: | |
# List all files in the downloads directory | |
files = list_repo_files( | |
repo_id=repo_id, | |
repo_type="space", | |
revision="main" | |
) | |
# Filter for video files in the downloads directory | |
video_files = [f for f in files if f.startswith('downloads/') and f.endswith('.mp4')] | |
if not video_files: | |
print("No video files found in the downloads directory.") | |
return | |
print(f"Found {len(video_files)} video files to download.") | |
# Download each video file | |
for video_file in video_files: | |
print(f"\nDownloading {video_file}...") | |
hf_hub_download( | |
repo_id=repo_id, | |
filename=video_file, | |
repo_type="space", | |
local_dir=local_dir | |
) | |
print(f"Downloaded {video_file} successfully!") | |
print("\nAll video files downloaded successfully!") | |
except Exception as e: | |
print(f"An error occurred while downloading video files: {str(e)}") | |
raise | |
def download_and_extract_precomputed(): | |
""" | |
Download and extract precomputed results from Hugging Face. | |
Downloads both intermediate_results.tar.xz and outputs.tar.xz, | |
and downloads all video files from the downloads directory. | |
""" | |
repo_id = "abdullahyang/NonisotropicSkeletonDiffusion_PrecomputedResults" | |
# Download and extract intermediate_results.tar.xz | |
download_and_extract_file( | |
repo_id=repo_id, | |
filename="intermediate_results.tar.xz", | |
local_dir="." | |
) | |
# Download and extract outputs.tar.xz | |
download_and_extract_file( | |
repo_id=repo_id, | |
filename="outputs.tar.xz", | |
local_dir="." | |
) | |
# Download video files | |
download_video_files( | |
repo_id=repo_id, | |
local_dir="." | |
) | |
print("All downloads and extractions completed successfully!") | |
if __name__ == "__main__": | |
download_and_extract_precomputed() |