SkeletonDiffusion/SkeletonDiffusion_Demo

1import os2from huggingface_hub import hf_hub_download, list_repo_files3import tarfile4import shutil5from tqdm import tqdm6 7def download_and_extract_file(repo_id, filename, local_dir):8    """9    Download and extract a file from Hugging Face.10    11    Args:12        repo_id (str): The repository ID on Hugging Face13        filename (str): The name of the file to download14        local_dir (str): The directory to extract files to15    """16    print(f"Starting download of {filename}...")17    18    # Create directory if it doesn't exist19    os.makedirs(local_dir, exist_ok=True)20    21    try:22        # Download the file using hf_hub_download23        file_path = hf_hub_download(24            repo_id=repo_id,25            filename=filename,26            repo_type="space",27            local_dir=local_dir28        )29        30        print(f"Download completed. File saved to: {file_path}")31        32        # Extract the tar.xz file with progress bar33        print(f"Extracting {filename}...")34        with tarfile.open(file_path, 'r:xz') as tar:35            # Get total number of files36            total_files = len(tar.getmembers())37            38            # Extract with progress bar39            for member in tqdm(tar.getmembers(), total=total_files, desc=f"Extracting {filename}"):40                tar.extract(member, path=local_dir)41        42        print(f"Extraction of {filename} completed successfully!")43        44        # Clean up the downloaded archive45        os.remove(file_path)46        print(f"Cleaned up downloaded archive {filename}.")47        48    except Exception as e:49        print(f"An error occurred while processing {filename}: {str(e)}")50        # Clean up in case of error51        if os.path.exists(file_path):52            os.remove(file_path)53        raise54 55def download_video_files(repo_id, local_dir):56    """57    Download all video files from the downloads directory.58    59    Args:60        repo_id (str): The repository ID on Hugging Face61        local_dir (str): The directory to save videos to62    """63    print("Starting download of video files...")64    65    # Create downloads directory if it doesn't exist66    os.makedirs(local_dir, exist_ok=True)67    68    try:69        # List all files in the downloads directory70        files = list_repo_files(71            repo_id=repo_id,72            repo_type="space",73            revision="main"74        )75        76        # Filter for video files in the downloads directory77        video_files = [f for f in files if f.startswith('downloads/') and f.endswith('.mp4')]78        79        if not video_files:80            print("No video files found in the downloads directory.")81            return82        83        print(f"Found {len(video_files)} video files to download.")84        85        # Download each video file86        for video_file in video_files:87            print(f"\nDownloading {video_file}...")88            hf_hub_download(89                repo_id=repo_id,90                filename=video_file,91                repo_type="space",92                local_dir=local_dir93            )94            print(f"Downloaded {video_file} successfully!")95        96        print("\nAll video files downloaded successfully!")97        98    except Exception as e:99        print(f"An error occurred while downloading video files: {str(e)}")100        raise101 102def download_and_extract_precomputed():103    """104    Download and extract precomputed results from Hugging Face.105    Downloads both intermediate_results.tar.xz and outputs.tar.xz,106    and downloads all video files from the downloads directory.107    """108    repo_id = "abdullahyang/NonisotropicSkeletonDiffusion_PrecomputedResults"109    110    # Download and extract intermediate_results.tar.xz111    download_and_extract_file(112        repo_id=repo_id,113        filename="intermediate_results.tar.xz",114        local_dir="."115    )116    117    # Download and extract outputs.tar.xz118    download_and_extract_file(119        repo_id=repo_id,120        filename="outputs.tar.xz",121        local_dir="."122    )123    124    # Download video files125    download_video_files(126        repo_id=repo_id,127        local_dir="."128    )129    130    print("All downloads and extractions completed successfully!")131 132if __name__ == "__main__":133    download_and_extract_precomputed()