60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
# layerstyle advance
|
|
|
|
import os.path
|
|
from pathlib import Path
|
|
import torch
|
|
from PIL import Image
|
|
import math
|
|
from torchvision.transforms import ToPILImage
|
|
import folder_paths
|
|
from .imagefunc import files_for_uform_gen2_qwen, StopOnTokens, UformGen2QwenChat, clear_memory, log
|
|
|
|
NODE_NAME = "QWenImage2Prompt"
|
|
# Example of integrating UformGen2QwenChat into a node-like structure
|
|
class QWenImage2Prompt:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"question": ("STRING", {"multiline": False, "default": "describe this image",},),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "uform_gen2_qwen_chat"
|
|
CATEGORY = '😺dzNodes/LayerUtility/Prompt'
|
|
|
|
def uform_gen2_qwen_chat(self, image, question):
|
|
chat_model = UformGen2QwenChat()
|
|
history = [] # Example empty history
|
|
pil_image = ToPILImage()(image[0].permute(2, 0, 1))
|
|
width, height = pil_image.size
|
|
ratio = width / height
|
|
if width * height > 1024 * 1024:
|
|
target_width = math.sqrt(ratio * 1024 * 1024)
|
|
target_height = target_width / ratio
|
|
target_width = int(target_width)
|
|
target_height = int(target_height)
|
|
pil_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
|
|
temp_path = files_for_uform_gen2_qwen / "temp.png"
|
|
pil_image.save(temp_path)
|
|
question = f"{question} but output no more then 80 words."
|
|
response = chat_model.chat_response(question, history, temp_path)
|
|
|
|
# Cleanup
|
|
del chat_model
|
|
clear_memory()
|
|
ret_text = response.split("assistant\n", 1)[1]
|
|
log(f"{NODE_NAME} Processed, Question: {question}, Response: {ret_text} ", message_type='finish')
|
|
return (ret_text, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayerUtility: QWenImage2Prompt": QWenImage2Prompt
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayerUtility: QWenImage2Prompt": "LayerUtility: QWenImage2Prompt(Advance)"
|
|
} |