From 18eb5498eee39addb11ca77a55fd67f32bc12ec1 Mon Sep 17 00:00:00 2001
From: Niels <niels.rogge1@gmail.com>
Date: Sun, 2 Jun 2024 11:43:29 +0200
Subject: [PATCH] Add metadata

---
 README.md                           | 12 ++++++++++++
 ultralytics/engine/model.py         |  4 +---
 ultralytics/models/yolov10/model.py |  4 +++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 0e8bfec8..f348a062 100644
--- a/README.md
+++ b/README.md
@@ -104,6 +104,18 @@ model.train(data='coco.yaml', epochs=500, batch=256, imgsz=640)
 # model.push_to_hub("reponame")
 ```
 
+## Push to hub to 🤗
+
+Optionally, you can push your fine-tuned model to the [Hugging Face hub](https://huggingface.co/) as a public or private model:
+
+```python
+# let's say you have fine-tuned a model for crop detection
+model.push_to_hub("<your-hf-username-or-organization/yolov10-finetuned-crop-detection")
+
+# you can also pass `private=True` if you don't want everyone to see your model
+model.push_to_hub("<your-hf-username-or-organization/yolov10-finetuned-crop-detection", private=True)
+```
+
 ## Prediction
 Note that a smaller confidence threshold can be set to detect smaller objects or objects in the distance. Please refer to [here](https://github.com/THU-MIG/yolov10/issues/136) for details.
 ```
diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py
index 4339dbd5..ef5c93c0 100644
--- a/ultralytics/engine/model.py
+++ b/ultralytics/engine/model.py
@@ -13,10 +13,8 @@ from ultralytics.hub.utils import HUB_WEB_ROOT
 from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load
 from ultralytics.utils import ASSETS, DEFAULT_CFG_DICT, LOGGER, RANK, SETTINGS, callbacks, checks, emojis, yaml_load
 
-from huggingface_hub import PyTorchModelHubMixin
 
-
-class Model(nn.Module, PyTorchModelHubMixin):
+class Model(nn.Module):
     """
     A base class for implementing YOLO models, unifying APIs across different model types.
 
diff --git a/ultralytics/models/yolov10/model.py b/ultralytics/models/yolov10/model.py
index a8881395..6e8eae42 100644
--- a/ultralytics/models/yolov10/model.py
+++ b/ultralytics/models/yolov10/model.py
@@ -4,7 +4,9 @@ from .val import YOLOv10DetectionValidator
 from .predict import YOLOv10DetectionPredictor
 from .train import YOLOv10DetectionTrainer
 
-class YOLOv10(Model):
+from huggingface_hub import PyTorchModelHubMixin
+
+class YOLOv10(Model, PyTorchModelHubMixin, library_name="ultralytics", repo_url="https://github.com/THU-MIG/yolov10", tags=["object-detection", "yolov10"]):
 
     def __init__(self, model="yolov10n.pt", task=None, verbose=False):
         super().__init__(model=model, task=task, verbose=verbose)