From 8e839632776a535806c80e1fd41fee8163365e35 Mon Sep 17 00:00:00 2001 From: Zernel Date: Tue, 31 Dec 2024 13:57:47 +0800 Subject: [PATCH] Fix bug in Moondream node --- moondream.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/moondream.py b/moondream.py index 58d1f77..a845945 100644 --- a/moondream.py +++ b/moondream.py @@ -14,13 +14,16 @@ class Moondream: } RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("Image Name",) + RETURN_NAMES = ("Image Description",) - FUNCTION = "get_image_name" + FUNCTION = "describe_image" CATEGORY = "image/utils" - def get_image_name(self, image): + def describe_image(self, image): + # change tensor to PIL Image + pil_image = Image.fromarray((image[0] * 255).numpy().astype('uint8')) + model_id = "vikhyatk/moondream2" revision = "2024-08-26" model = AutoModelForCausalLM.from_pretrained( @@ -28,5 +31,5 @@ class Moondream: ) tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision) - enc_image = model.encode_image(image) + enc_image = model.encode_image(pil_image) return (model.answer_question(enc_image, "Describe this image.", tokenizer),)