81 lines
2.4 KiB
Python
Executable File
81 lines
2.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
from huggingface_hub import HfApi
|
|
import argparse
|
|
import sys
|
|
|
|
def add_model_to_collection(collection_slug, model_id, note=""):
|
|
"""
|
|
Add a model to an existing collection
|
|
|
|
Args:
|
|
collection_slug: The slug of the collection (e.g., "username/collection-name-12345")
|
|
model_id: The model repository ID (e.g., "username/model-name")
|
|
note: Optional note about the model
|
|
|
|
Returns:
|
|
True if successful, False if failed
|
|
"""
|
|
|
|
# Initialize API
|
|
api = HfApi()
|
|
|
|
try:
|
|
user_info = api.whoami()
|
|
print(f"✅ Authenticated as: {user_info['name']}")
|
|
|
|
# Verify the model exists
|
|
print(f"🔍 Checking if model exists: {model_id}")
|
|
try:
|
|
model_info = api.model_info(model_id)
|
|
except Exception as e:
|
|
print(f"❌ Model not found or not accessible: {model_id}")
|
|
print(f"Error: {e}")
|
|
return False
|
|
|
|
print(f"📚 Adding model to collection...")
|
|
api.add_collection_item(
|
|
collection_slug=collection_slug,
|
|
item_id=model_id,
|
|
item_type="model",
|
|
note=note
|
|
)
|
|
|
|
print(f"✅ Model added to collection successfully!")
|
|
print(f"🔗 Collection URL: https://huggingface.co/collections/{collection_slug}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error adding model to collection: {e}")
|
|
return False
|
|
|
|
def main():
|
|
# This script requires that the environment variable HF_TOKEN is set with your
|
|
# Hugging Face API token.
|
|
api = HfApi()
|
|
|
|
parser = argparse.ArgumentParser(description='Add model to a Huggingface Collection')
|
|
parser.add_argument('--collection', '-c', help='The collection slug username/collection-hash', required=True)
|
|
parser.add_argument('--model', '-m', help='The model to add to the Collection', required=True)
|
|
parser.add_argument('--note', '-n', help='An optional note/description', required=False)
|
|
args = parser.parse_args()
|
|
|
|
collection = args.collection
|
|
model = args.model
|
|
note = args.note
|
|
|
|
success = add_model_to_collection(
|
|
collection_slug=collection,
|
|
model_id=model,
|
|
note=note
|
|
)
|
|
|
|
if success:
|
|
print("\n🎉 Model added successfully!")
|
|
else:
|
|
print("\n❌ Failed to add model to collection")
|
|
sys.exit(1)
|
|
if __name__ == "__main__":
|
|
main()
|