diff --git a/__init__.py b/__init__.py index 1fad1b4..10d2909 100644 --- a/__init__.py +++ b/__init__.py @@ -1,12 +1,14 @@ from .upload_images import UploadImages from .classifiers import * - +from .moondream import Moondream NODE_CLASS_MAPPINGS = { "UploadImagesToInstantStudio": UploadImages, 'HuggingFace Classify': HFClassify, + 'Moondream': Moondream, } NODE_DISPLAY_NAME_MAPPINGS = { - "UploadImagesToInstantStudio": "Upload Images to Instant Studio" + "UploadImagesToInstantStudio": "Upload Images to Instant Studio", + "Moondream": "Vision LM to analyze images" } diff --git a/moondream.py b/moondream.py new file mode 100644 index 0000000..58d1f77 --- /dev/null +++ b/moondream.py @@ -0,0 +1,32 @@ +from transformers import AutoModelForCausalLM, AutoTokenizer +from PIL import Image + +class Moondream: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("Image Name",) + + FUNCTION = "get_image_name" + + CATEGORY = "image/utils" + + def get_image_name(self, image): + model_id = "vikhyatk/moondream2" + revision = "2024-08-26" + model = AutoModelForCausalLM.from_pretrained( + model_id, trust_remote_code=True, revision=revision + ) + tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision) + + enc_image = model.encode_image(image) + return (model.answer_question(enc_image, "Describe this image.", tokenizer),) diff --git a/requirements.txt b/requirements.txt index d6a5666..94c9497 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ numpy Pillow Requests +flash_attn +transformers +einops \ No newline at end of file