1#!/usr/bin/env python3
  2
  3from huggingface_hub import HfApi
  4import argparse
  5import os
  6import sys
  7
  8
  9def create_collection(title, description, private=False, namespace=None, return_slug=False):
 10    """
 11    Create a new collection on Hugging Face
 12
 13    Args:
 14        title: Collection title
 15        description: Collection description
 16        private: Whether the collection should be private (default: False)
 17        namespace: Optional namespace (defaults to your username)
 18
 19    Returns:
 20        Collection object if successful, None if failed
 21    """
 22
 23    # Check if HF_TOKEN is available
 24    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
 25    if not token:
 26        print("❌ No HF_TOKEN or HUGGINGFACE_HUB_TOKEN found in environment variables")
 27        print("Please set your Hugging Face token as an environment variable")
 28        return None
 29
 30    # Initialize API
 31    api = HfApi()
 32
 33    try:
 34        # Test authentication first
 35        user_info = api.whoami()
 36        if not return_slug:
 37            print(f"✅ Authenticated as: {user_info['name']}")
 38
 39        # Create the collection
 40        if not return_slug:
 41            print(f"📚 Creating collection: '{title}'...")
 42        collection = api.create_collection(
 43            title=title,
 44            description=description,
 45            private=private,
 46            namespace=namespace
 47        )
 48
 49        if not return_slug:
 50            print(f"✅ Collection created successfully!")
 51            print(f"📋 Collection slug: {collection.slug}")
 52            print(f"🔗 Collection URL: https://huggingface.co/collections/{collection.slug}")
 53
 54        return collection
 55
 56    except Exception as e:
 57        print(f"❌ Error creating collection: {e}")
 58        return None
 59
 60def main():
 61    # This script requires that the environment variable HF_TOKEN is set with your
 62    # Hugging Face API token.
 63    api = HfApi()
 64
 65    parser = argparse.ArgumentParser(description='Create a Huggingface Collection')
 66    parser.add_argument('--name', '-n', help='The name/title of the Collection', required=True)
 67    parser.add_argument('--description', '-d', help='The description for the Collection', required=True)
 68    parser.add_argument('--namespace', '-ns', help='The namespace to add the Collection to', required=True)
 69    parser.add_argument('--private', '-p', help='Create a private Collection', action='store_true')  # Fixed
 70    parser.add_argument('--return-slug', '-s', help='Only output the collection slug', action='store_true')  # Fixed
 71
 72    args = parser.parse_args()
 73
 74    name = args.name
 75    description = args.description
 76    private = args.private
 77    namespace = args.namespace
 78    return_slug = args.return_slug
 79
 80    if not return_slug:
 81        print("🚀 Creating Hugging Face Collection")
 82        print(f"Title: {name}")
 83        print(f"Description: {description}")
 84        print(f"Namespace: {namespace}")
 85        print(f"Private: {private}")
 86
 87    collection = create_collection(
 88        title=name,
 89        description=description,
 90        private=private,
 91        namespace=namespace,
 92        return_slug=return_slug
 93    )
 94
 95    if collection:
 96        if return_slug:
 97            print(collection.slug)
 98        else:
 99            print("\n🎉 Collection created successfully!")
100            print(f"Use this slug to add models: {collection.slug}")
101    else:
102        print("\n❌ Failed to create collection")
103        sys.exit(1)
104
105if __name__ == "__main__":
106    main()