SkeletonDiffusion_Demo / download_precomputed.py
Jialin Yang
Initial release on Huggingface Spaces with Gradio UI
352b049
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()