Update image_captioner.py

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 (</think>)
+            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",
 }