diff --git a/README.md b/README.md index 1f625ef45..eeeb1d4f1 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,8 @@ Currently implemented services: - Transport - Daily - Local (in progress, intended as a quick start example service) +- Vision + - Moondream If you'd like to [implement a service]((https://github.com/daily-co/daily-ai-sdk/tree/main/src/dailyai/services)), we welcome PRs! Our goal is to support lots of services in all of the above categories, plus new categories (like real-time video) as they emerge. @@ -63,7 +65,7 @@ pip install "dailyai[option,...]" Your project may or may not need these, so they're made available as optional requirements. Here is a list: -- **AI services**: `anthropic`, `azure`, `fal`, `openai`, `playht`, `silero`, `whisper` +- **AI services**: `anthropic`, `azure`, `fal`, `moondream`, `openai`, `playht`, `silero`, `whisper` - **Transports**: `daily`, `local`, `websocket` ## Code examples diff --git a/pyproject.toml b/pyproject.toml index 8baf090ae..c42452d4d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ daily = [ "daily-python~=0.7.0" ] examples = [ "python-dotenv~=1.0.0", "flask~=3.0.0", "flask_cors~=4.0.0" ] fal = [ "fal~=0.12.0" ] local = [ "pyaudio~=0.2.0" ] +moondream = [ "einops~=0.7.0", "timm~=0.9.0", "transformers~=4.39.0" ] openai = [ "openai~=1.14.0" ] playht = [ "pyht~=0.0.26" ] silero = [ "torch~=2.2.0", "torchaudio~=2.2.0" ] diff --git a/src/dailyai/services/moondream_ai_service.py b/src/dailyai/services/moondream_ai_service.py new file mode 100644 index 000000000..2784a669e --- /dev/null +++ b/src/dailyai/services/moondream_ai_service.py @@ -0,0 +1,52 @@ +from dailyai.pipeline.frames import ImageFrame +from dailyai.services.ai_services import VisionService + +from PIL import Image + +from transformers import AutoModelForCausalLM, AutoTokenizer + +import torch + + +def detect_device(): + """ + Detects the appropriate device to run on, and return the device and dtype. + """ + if torch.cuda.is_available(): + return torch.device("cuda"), torch.float16 + elif torch.backends.mps.is_available(): + return torch.device("mps"), torch.float16 + else: + return torch.device("cpu"), torch.float32 + + +class MoondreamService(VisionService): + def __init__( + self, + model_id="vikhyatk/moondream2", + revision="2024-04-02", + device=None + ): + super().__init__() + + if not device: + device, dtype = detect_device() + else: + device = torch.device("cpu") + dtype = torch.float32 + + self._tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision) + + self._model = AutoModelForCausalLM.from_pretrained( + model_id, trust_remote_code=True, revision=revision + ).to(device=device, dtype=dtype) + self._model.eval() + + async def run_vision(self, describe_text: str, frame: ImageFrame) -> str: + image = Image.frombytes("RGB", (frame.size[0], frame.size[1]), frame.image) + image_embeds = self._model.encode_image(image) + description = self._model.answer_question( + image_embeds=image_embeds, + question=describe_text, + tokenizer=self._tokenizer) + return description