diff --git a/pyproject.toml b/pyproject.toml index 04ead153..9a7135ef 100644 --- a/pyproject.toml +++ b/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 "] @@ -72,6 +72,9 @@ qdrant-client = "*" sentence-transformers = "*" peft = "*" psutil = "*" +ultralytics = "*" +timm = "*" + [tool.poetry.group.lint.dependencies] diff --git a/requirements.txt b/requirements.txt index 29b6cbe5..3c14d025 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file +psutil +ultralytics \ No newline at end of file diff --git a/swarms/models/__init__.py b/swarms/models/__init__.py index 8ce74b27..0b1992eb 100644 --- a/swarms/models/__init__.py +++ b/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", ] diff --git a/swarms/models/timm.py b/swarms/models/timm.py index dc1e41b9..de0484f2 100644 --- a/swarms/models/timm.py +++ b/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 diff --git a/swarms/models/ultralytics_model.py b/swarms/models/ultralytics_model.py new file mode 100644 index 00000000..2df548f4 --- /dev/null +++ b/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)