130 lines
4.2 KiB
Python
130 lines
4.2 KiB
Python
import torch
|
|
import random
|
|
import re
|
|
|
|
# Qwen3用チャットテンプレート
|
|
QWEN3_CHAT_TEMPLATE = "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n"
|
|
|
|
# デフォルトシステムプロンプト
|
|
DEFAULT_SYSTEM_PROMPT = "You are a professional image generation prompt expert. Output in English."
|
|
|
|
|
|
def build_qwen3_prompt(system_prompt: str, user_prompt: str, enable_thinking: bool = True) -> str:
|
|
"""Qwen3フォーマットでプロンプトを構築。
|
|
Non-thinkingモード時は /no_think ソフトスイッチをユーザーメッセージに付与する。
|
|
"""
|
|
if not enable_thinking:
|
|
user_prompt = user_prompt + "\n/no_think"
|
|
return QWEN3_CHAT_TEMPLATE.format(system=system_prompt, user=user_prompt)
|
|
|
|
|
|
class TextGenerateQwen3Prompt:
|
|
CATEGORY = "LLM"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("generated_text",)
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"system_prompt": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": DEFAULT_SYSTEM_PROMPT,
|
|
},
|
|
),
|
|
"max_length": ("INT", {"default": 256, "min": 1, "max": 4096}),
|
|
"do_sample": ("BOOLEAN", {"default": True}),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01},
|
|
),
|
|
"top_k": ("INT", {"default": 64, "min": 0, "max": 1000}),
|
|
"top_p": (
|
|
"FLOAT",
|
|
{"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"min_p": (
|
|
"FLOAT",
|
|
{"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"repetition_penalty": (
|
|
"FLOAT",
|
|
{"default": 1.05, "min": 1.0, "max": 2.0, "step": 0.01},
|
|
),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"enable_thinking": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
clip,
|
|
prompt,
|
|
system_prompt,
|
|
max_length,
|
|
do_sample,
|
|
temperature,
|
|
top_k,
|
|
top_p,
|
|
min_p,
|
|
repetition_penalty,
|
|
seed,
|
|
enable_thinking,
|
|
):
|
|
# シード設定
|
|
torch.manual_seed(seed)
|
|
random.seed(seed)
|
|
|
|
# Qwen3フォーマットでプロンプト構築 (enable_thinkingをテンプレートで制御)
|
|
full_prompt = build_qwen3_prompt(system_prompt, prompt, enable_thinking)
|
|
|
|
# トークナイズ
|
|
tokens = clip.tokenize(full_prompt)
|
|
|
|
# 生成パラメータ設定
|
|
gen_kwargs = {
|
|
"max_length": max_length,
|
|
"do_sample": do_sample,
|
|
"temperature": temperature,
|
|
"top_k": top_k,
|
|
"top_p": top_p,
|
|
"min_p": min_p,
|
|
"repetition_penalty": repetition_penalty,
|
|
"seed": seed,
|
|
}
|
|
|
|
# テキスト生成
|
|
output = clip.generate(tokens, **gen_kwargs)
|
|
|
|
# デコード
|
|
if hasattr(clip, "decode"):
|
|
generated_text = clip.decode(output)
|
|
else:
|
|
generated_text = output
|
|
|
|
# プロンプト部分を除去
|
|
if isinstance(generated_text, str) and generated_text.startswith(full_prompt):
|
|
generated_text = generated_text[len(full_prompt):]
|
|
|
|
# <think>...</think> ブロックを除去
|
|
generated_text = re.sub(r"<think>.*?</think>", "", generated_text, flags=re.DOTALL)
|
|
|
|
# 特殊トークンのクリーンアップ
|
|
for marker in ["<|im_start|>", "<|im_end|>"]:
|
|
generated_text = generated_text.replace(marker, "")
|
|
|
|
return (generated_text.strip(),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TextGenerateQwen3Prompt": TextGenerateQwen3Prompt,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TextGenerateQwen3Prompt": "Text Generate Qwen3 Prompt",
|
|
} |