publish_to_hf.py

"""
Upload Byte Dream to Hugging Face Hub
Using the new push_to_hub API from generator
"""

import os
import sys
from pathlib import Path


def main():
    # Get token from command line argument or prompt
    if len(sys.argv) > 1:
        token = sys.argv[1]
    else:
        token = input("Enter your Hugging Face token: ")
    
    # Get repo ID
    if len(sys.argv) > 2:
        repo_id = sys.argv[2]
    else:
        repo_id = input("Enter repository ID (e.g., username/ByteDream): ")
    
    # Check if model exists, if not train first
    model_path = Path("./models/bytedream")
    if not model_path.exists():
        print("\n⚠ Model directory not found!")
        print("Please train the model first using: python train.py")
        print("Or download pretrained weights.")
        sys.exit(1)
    
    print("\n" + "="*60)
    print("Byte Dream - Upload to Hugging Face Hub")
    print("="*60)
    
    try:
        # Import generator
        from bytedream.generator import ByteDreamGenerator
        
        # Initialize generator with trained model
        print("\nLoading model...")
        generator = ByteDreamGenerator(
            model_path="./models/bytedream",
            config_path="config.yaml",
            device="cpu",
        )
        
        # Push to Hub
        print(f"\nUploading to {repo_id}...")
        generator.push_to_hub(
            repo_id=repo_id,
            token=token,
            private=False,
            commit_message="Upload Byte Dream model",
        )
        
        print("\n" + "="*60)
        print("✓ SUCCESS!")
        print("="*60)
        print(f"\n📦 Your model is available at:")
        print(f"https://huggingface.co/{repo_id}")
        print("\nTo use this model:")
        print(f"  python infer.py --prompt 'your prompt' --model '{repo_id}'")
        print("="*60)
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()