1#!/usr/bin/env python3
 2
 3from huggingface_hub import HfApi
 4import argparse
 5import sys
 6
 7def add_model_to_collection(collection_slug, model_id, note=""):
 8    """
 9    Add a model to an existing collection
10
11    Args:
12        collection_slug: The slug of the collection (e.g., "username/collection-name-12345")
13        model_id: The model repository ID (e.g., "username/model-name")
14        note: Optional note about the model
15
16    Returns:
17        True if successful, False if failed
18    """
19
20    # Initialize API
21    api = HfApi()
22
23    try:
24        user_info = api.whoami()
25        print(f"✅ Authenticated as: {user_info['name']}")
26
27        # Verify the model exists
28        print(f"🔍 Checking if model exists: {model_id}")
29        try:
30            model_info = api.model_info(model_id)
31        except Exception as e:
32            print(f"❌ Model not found or not accessible: {model_id}")
33            print(f"Error: {e}")
34            return False
35
36        print(f"📚 Adding model to collection...")
37        api.add_collection_item(
38            collection_slug=collection_slug,
39            item_id=model_id,
40            item_type="model",
41            note=note
42        )
43
44        print(f"✅ Model added to collection successfully!")
45        print(f"🔗 Collection URL: https://huggingface.co/collections/{collection_slug}")
46
47        return True
48
49    except Exception as e:
50        print(f"❌ Error adding model to collection: {e}")
51        return False
52
53def main():
54    # This script requires that the environment variable HF_TOKEN is set with your
55    # Hugging Face API token.
56    api = HfApi()
57
58    parser = argparse.ArgumentParser(description='Add model to a Huggingface Collection')
59    parser.add_argument('--collection', '-c', help='The collection slug username/collection-hash', required=True)
60    parser.add_argument('--model', '-m', help='The model to add to the Collection', required=True)
61    parser.add_argument('--note', '-n', help='An optional note/description', required=False)
62    args = parser.parse_args()
63
64    collection = args.collection
65    model = args.model
66    note = args.note
67
68    success = add_model_to_collection(
69        collection_slug=collection,
70        model_id=model,
71        note=note
72    )
73
74    if success:
75        print("\n🎉 Model added successfully!")
76    else:
77        print("\n❌ Failed to add model to collection")
78        sys.exit(1)
79if __name__ == "__main__":
80    main()