deploy_to_spaces.py

"""
Deploy Byte Dream to Hugging Face Spaces
Automated deployment script for creating and updating Spaces
"""

import argparse
from pathlib import Path
import shutil
import tempfile
from huggingface_hub import HfApi, create_repo
import os


def deploy_to_spaces(
    repo_id: str,
    space_title: str = "Byte Dream - AI Image Generator",
    token: str = None,
    private: bool = False,
    hardware: str = "cpu-basic",
):
    """
    Deploy Byte Dream to Hugging Face Spaces
    
    Args:
        repo_id: Space repository ID (username/SpaceName)
        space_title: Title for the Space
        token: Hugging Face API token
        private: Whether to make Space private
        hardware: Hardware tier (cpu-basic, cpu-upgrade, etc.)
    """
    api = HfApi()
    
    print("=" * 60)
    print("Byte Dream - Hugging Face Spaces Deployment")
    print("=" * 60)
    
    # Create space if it doesn't exist
    try:
        create_repo(
            repo_id=repo_id,
            token=token,
            repo_type="space",
            space_sdk="gradio",
            space_hardware=hardware,
            private=private,
            exist_ok=True,
        )
        print(f"✓ Space created/verified: {repo_id}")
    except Exception as e:
        print(f"Error creating space: {e}")
        return
    
    # Prepare files to upload
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = Path(tmp_dir)
        
        # Copy essential files
        files_to_copy = [
            "app.py",
            "config.yaml",
            "requirements.txt",
        ]
        
        for file in files_to_copy:
            src = Path(file)
            if src.exists():
                shutil.copy2(src, tmp_path / file)
                print(f"✓ Copied {file}")
        
        # Copy bytedream package
        bytedream_src = Path("bytedream")
        bytedream_dst = tmp_path / "bytedream"
        if bytedream_src.exists():
            shutil.copytree(bytedream_src, bytedream_dst)
            print("✓ Copied bytedream package")
        
        # Copy spaces app template
        spaces_app = Path("spaces_app.py")
        if spaces_app.exists():
            # Rename to README.md for the space
            shutil.copy2(spaces_app, tmp_path / "README.md")
            print("✓ Copied README.md")
        
        # Create .gitignore for space
        gitignore_content = """
__pycache__/
*.pyc
.env
*.log
outputs/
models/
        """
        (tmp_path / ".gitignore").write_text(gitignore_content)
        print("✓ Created .gitignore")
        
        # Upload to space
        print("\nUploading files to Hugging Face Spaces...")
        try:
            api.upload_folder(
                folder_path=str(tmp_path),
                repo_id=repo_id,
                repo_type="space",
                token=token,
                commit_message="Deploy Byte Dream application",
            )
            print("✓ Files uploaded successfully!")
            
            # Get space URL
            space_url = f"https://huggingface.co/spaces/{repo_id}"
            print(f"\n🚀 Your Space is deployed at:")
            print(f"{space_url}")
            print(f"\nNote: It may take a few minutes for the Space to build and become available.")
            print(f"Hardware: {hardware}")
            
        except Exception as e:
            print(f"Error uploading to space: {e}")
            raise


def update_space(repo_id: str, files: list, token: str = None):
    """
    Update specific files in an existing Space
    
    Args:
        repo_id: Space repository ID
        files: List of files to update
        token: Hugging Face API token
    """
    api = HfApi()
    
    print(f"Updating files in space: {repo_id}")
    
    for file_path in files:
        if Path(file_path).exists():
            api.upload_file(
                path_or_fileobj=file_path,
                path_in_repo=Path(file_path).name,
                repo_id=repo_id,
                repo_type="space",
                token=token,
            )
            print(f"✓ Updated {file_path}")
        else:
            print(f"⚠ File not found: {file_path}")


def main():
    parser = argparse.ArgumentParser(description="Deploy Byte Dream to Hugging Face Spaces")
    
    parser.add_argument(
        "--repo_id",
        type=str,
        required=True,
        help="Space repository ID (e.g., username/ByteDream-Space)"
    )
    
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="Hugging Face API token (optional, will use cached token if not provided)"
    )
    
    parser.add_argument(
        "--title",
        type=str,
        default="Byte Dream - AI Image Generator",
        help="Title for the Space"
    )
    
    parser.add_argument(
        "--private",
        action="store_true",
        help="Make the Space private"
    )
    
    parser.add_argument(
        "--hardware",
        type=str,
        default="cpu-basic",
        choices=["cpu-basic", "cpu-upgrade", "cpu-xl", "zero-a10g", "t4-small", "t4-medium", "a10g-small", "a10g-large"],
        help="Hardware tier for the Space"
    )
    
    parser.add_argument(
        "--update",
        nargs="+",
        help="Update specific files instead of full deployment"
    )
    
    args = parser.parse_args()
    
    if args.update:
        # Update specific files
        update_space(args.repo_id, args.update, args.token)
    else:
        # Full deployment
        deploy_to_spaces(
            repo_id=args.repo_id,
            space_title=args.title,
            token=args.token,
            private=args.private,
            hardware=args.hardware,
        )


if __name__ == "__main__":
    main()