diff --git a/image_captioner.py b/image_captioner.py index 23332ec..d1648c8 100644 --- a/image_captioner.py +++ b/image_captioner.py @@ -4,7 +4,7 @@ from PIL import Image import os import time from transformers import pipeline - +from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import ( AutoModelForImageTextToText, AutoProcessor @@ -112,10 +112,78 @@ class ImageCaptioner: +class Quen3Helper: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "your_prompt": ("STRING", {"multiline": True, "default": ""}), + "model_id": ("STRING", {"multiline": True, "default": ""}), + } + } + + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("quen3_response",) + FUNCTION = "generate_captions" + OUTPUT_NODE = True + CATEGORY = "utils" + def generate_response(self,your_prompt, model_id,): + # Load the model and processor + model_name = model_id + tokenizer = AutoTokenizer.from_pretrained(model_name) + model = AutoModelForCausalLM.from_pretrained( + model_name, + torch_dtype="auto", + device_map="auto" + ) + # prepare the model input + prompt =your_prompt or "Give me a short introduction to large language model." + messages = [ + {"role": "user", "content": prompt} + ] + text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=True # Switches between thinking and non-thinking modes. Default is True. + ) + model_inputs = tokenizer([text], return_tensors="pt").to(model.device) + + # conduct text completion + generated_ids = model.generate( + **model_inputs, + max_new_tokens=32768 + ) + output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() + + # parsing thinking content + try: + # rindex finding 151668 () + index = len(output_ids) - output_ids[::-1].index(151668) + except ValueError: + index = 0 + + thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n") + content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n") + + print("thinking content:", thinking_content) + print("content:", content) + caption_file_path = os.path.join("./", "quenresponse.txt") + with open(caption_file_path, "w") as caption_file: + caption_file.write(content) + + return content, + + + + NODE_CLASS_MAPPINGS = { - "ImageCaptioner": ImageCaptioner + "ImageCaptioner": ImageCaptioner, + "Quen3Helper": Quen3Helper, } NODE_DISPLAY_NAME_MAPPINGS = { "ImageCaptioner": "Image Captioner", + "Quen3Helper": "Quen3 Helper", }