import os
import tarfile
from huggingface_hub import HfApi

# Initialize Hugging Face API
api = HfApi()
# Define paths and filenames
source_dir = "/workspace/logs/codebooks/kmeans/stable_vae_16384"
output_filename = "vq_codebook.tar.gz"
print(f"Creating tar.gz archive from {source_dir}...")
# Create tar.gz archive
with tarfile.open(output_filename, "w:gz") as tar:
    tar.add(source_dir, arcname=os.path.basename(source_dir))
print("Archive created successfully")

# Upload to Hugging Face
print("Uploading to Hugging Face Hub...")
api.upload_file(
    path_or_fileobj=output_filename,
    path_in_repo=output_filename,
    repo_id="seungheondoh/model_temp",  # Replace with your actual repo name
    repo_type="model"
)

print("Upload completed successfully")

# Clean up the local tar.gz file
os.remove(output_filename)
print("Local archive cleaned up")