import os
from huggingface_hub import snapshot_download

def download_litert_model():
    repo_id = "litert-community/Z-Image-Turbo-LiteRT"
    
    # Define where you want to save the models locally
    local_dir = "./Z-Image-Turbo-LiteRT"
    
    print(f"🚀 Starting download of '{repo_id}'...")
    print(f"📂 Files will be saved to: {os.path.abspath(local_dir)}")
    print("⚠️  Note: This repository is > 6 GB. This may take a few minutes depending on your network connection.\n")
    
    try:
        # snapshot_download grabs the whole repo, including all split .tflite chunks
        downloaded_path = snapshot_download(
            repo_id=repo_id,
            local_dir=local_dir,
            local_dir_use_symlinks=False,  # Copies actual files directly into the folder
            resume_download=True           # Allows resuming if the connection breaks
        )
        
        print("\n✅ Download complete!")
        print(f"✨ All model chunks are now ready in: {downloaded_path}")
        
        # List the downloaded files to verify
        print("\nDownloaded files:")
        for file in sorted(os.listdir(downloaded_path)):
            if file.endswith(".tflite"):
                size_mb = os.path.getsize(os.path.join(downloaded_path, file)) / (1024 * 1024)
                print(f" └── {file} ({size_mb:.2f} MB)")
                
    except Exception as e:
        print(f"\n❌ An error occurred during download: {e}")

if __name__ == "__main__":
    download_litert_model()