hf-add-model-to-collection.py 2.4 KB

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