[FEATS][UltralyticsModel] [TimmModel]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.6.7"
+version = "3.6.8"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
@@ -72,6 +72,9 @@ qdrant-client = "*"
 sentence-transformers = "*"
 peft = "*"
 psutil = "*"
+ultralytics = "*"
+timm = "*"
+
 
 
 [tool.poetry.group.lint.dependencies]

requirements.txt

@@ -45,7 +45,7 @@ openai==0.28.0
 opencv-python==4.7.0.72
 prettytable==3.9.0
 safetensors==0.3.3
-timm==0.6.13
+timm
 torchmetrics
 webdataset
 marshmallow==3.19.0
@@ -59,4 +59,5 @@ mkdocs-material
 mkdocs-glightbox
 pre-commit==3.2.2
 peft
-psutil
+psutil
+ultralytics

swarms/models/__init__.py

@@ -40,6 +40,8 @@ from swarms.models.openai_tts import OpenAITTS  # noqa: E402
 from swarms.models.gemini import Gemini  # noqa: E402
 from swarms.models.gigabind import Gigabind  # noqa: E402
 from swarms.models.zeroscope import ZeroscopeTTV  # noqa: E402
+from swarms.models.timm import TimmModel  # noqa: E402
+from swarms.models.ultralytics_model import UltralyticsModel  # noqa: E402
 
 
 # from swarms.models.dalle3 import Dalle3
@@ -96,4 +98,6 @@ __all__ = [
     # "ModelScopePipeline",
     # "ModelScopeAutoModel",
     "TogetherLLM",
+    "TimmModel",
+    "UltralyticsModel",
 ]

swarms/models/timm.py

@@ -26,8 +26,14 @@ class TimmModel(BaseMultiModalModel):
     """
 
     def __init__(
-        self, model_name: str, pretrained: bool, in_chans: int
+        self,
+        model_name: str,
+        pretrained: bool,
+        in_chans: int,
+        *args,
+        **kwargs,
     ):
+        super().__init__(*args, **kwargs)
         self.model_name = model_name
         self.pretrained = pretrained
         self.in_chans = in_chans

swarms/models/ultralytics_model.py

@@ -0,0 +1,32 @@
+from swarms.models.base_multimodal_model import BaseMultiModalModel
+from ultralytics import YOLO
+
+
+class UltralyticsModel(BaseMultiModalModel):
+    def __init__(self, model_name: str, *args, **kwargs):
+        """
+        Initializes an instance of the Ultralytics model.
+
+        Args:
+            model_name (str): The name of the model.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
+        super().__init__(*args, **kwargs)
+        self.model_name = model_name
+
+        self.model = YOLO(model_name, *args, **kwargs)
+
+    def __call__(self, task: str, *args, **kwargs):
+        """
+        Calls the Ultralytics model.
+
+        Args:
+            task (str): The task to perform.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            The result of the model call.
+        """
+        return self.model(task, *args, **kwargs)