comfyui-mixlab-nodes/nodes/TextGenerateNode.py

from transformers import pipeline, set_seed,AutoTokenizer, AutoModelForSeq2SeqLM
import random
import re

import os,sys
import folder_paths

# from PIL import Image
import importlib.util

import comfy.utils
# import numpy as np
import torch
import random
from lark import Lark, Transformer, v_args


global _available
_available=True

def get_text_generator_path():
    try:
        return folder_paths.get_folder_paths('prompt_generator')[0]
    except:
        return os.path.join(folder_paths.models_dir, "prompt_generator")

prompt_generator=get_text_generator_path()

text_generator_model_path=os.path.join(prompt_generator, "text2image-prompt-generator")
if not os.path.exists(text_generator_model_path):
    print(f"## text_generator_model not found: {text_generator_model_path}, pls download from https://huggingface.co/succinctly/text2image-prompt-generator/tree/main")
    text_generator_model_path='succinctly/text2image-prompt-generator'

zh_en_model_path=os.path.join(prompt_generator, "opus-mt-zh-en")
if not os.path.exists(zh_en_model_path):
    print(f"## zh_en_model not found: {zh_en_model_path}, pls download from https://huggingface.co/Helsinki-NLP/opus-mt-zh-en/tree/main")
    zh_en_model_path='Helsinki-NLP/opus-mt-zh-en'


def is_installed(package):
    try:
        spec = importlib.util.find_spec(package)
    except ModuleNotFoundError:
        return False
    return spec is not None


try:
    if is_installed('sentencepiece')==False:
        import subprocess

        # 安装
        print('#pip install sentencepiece')

        result = subprocess.run([sys.executable, '-s', '-m', 'pip', 'install', 'sentencepiece'], capture_output=True, text=True)

        #检查命令执行结果
        if result.returncode == 0 and is_installed('sentencepiece'):
            print("#install success")
            _available=True
        else:
            print("#install error")
            _available=False
        
    else:
        _available=True

except:
    _available=False



def translate(text):
    global text_pipe,zh_en_model,zh_en_tokenizer
    
    if zh_en_model==None:
        zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
        zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path,padding=True, truncation=True)
    
    zh_en_model.to("cuda" if torch.cuda.is_available() else "cpu")
    with torch.no_grad():
        encoded = zh_en_tokenizer([text], return_tensors="pt")
        encoded.to(zh_en_model.device)
        sequences = zh_en_model.generate(**encoded)
        return zh_en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]

# input = "青春不能回头，所以青春没有终点。 ——《火影忍者》"
# print(input, translate(input))


def text_generate(text_pipe,input,seed=None):
    
    if seed==None:
        seed = random.randint(100, 1000000)
    
    set_seed(seed)

    for count in range(6):    
        sequences = text_pipe(input, max_length=random.randint(60, 90), num_return_sequences=8)
        list = []
        for sequence in sequences:
            line = sequence['generated_text'].strip()
            if line != input and len(line) > (len(input) + 4) and line.endswith((":", "-", "—")) is False:
                list.append(line)

        result = "\n".join(list)
        result = re.sub('[^ ]+\.[^ ]+','', result)
        result = result.replace("<", "").replace(">", "")
        if result != "":
            return result
        if count == 5:
            return result

# input = "Youth can't turn back, so there's no end to youth."
# print(input, text_generate(input))

    
import re

def correct_prompt_syntax(prompt=""):

    # print("input prompt",prompt)
    corrected_elements = []
    # 处理成统一的英文标点
    prompt = prompt.replace('（', '(').replace('）', ')').replace('，', ',').replace(';', ',').replace('。', '.').replace('：',':')
    # 删除多余的空格
    prompt = re.sub(r'\s+', ' ', prompt).strip()
    prompt = prompt.replace("< ","<").replace(" >",">").replace("( ","(").replace(" )",")").replace("[ ","[").replace(' ]',']')

    # 分词
    prompt_elements = prompt.split(',')

    def balance_brackets(element, open_bracket, close_bracket):
        open_brackets_count = element.count(open_bracket)
        close_brackets_count = element.count(close_bracket)
        return element + close_bracket * (open_brackets_count - close_brackets_count)

    for element in prompt_elements:
        element = element.strip()

        # 处理空元素
        if not element:
            continue

        # 检查并处理圆括号、方括号、尖括号
        if element[0] in '([':
            corrected_element = balance_brackets(element, '(', ')') if element[0] == '(' else balance_brackets(element, '[', ']')
        elif element[0] == '<':
            corrected_element = balance_brackets(element, '<', '>')
        else:
            # 删除开头的右括号或右方括号
            corrected_element = element.lstrip(')]')

        corrected_elements.append(corrected_element)

    # 重组修正后的prompt
    return  ','.join(corrected_elements)


# # 示例使用
# test_prompt = "((middle-century castles)), [forsaken: 0.8], (mystery dragons: 1.3, mist forests, sunsets, quiet; (((dummy)), [fisting city: 0.5] background, radiant, soft and flavoured,] promising mountains, ((starry: 1.6), [[crowds], [middle-century castle: urban landscapes of the future: 0.5], [yellow: bright sun: 0.7], overlooking"
# corrected_prompt = correct_prompt_syntax(test_prompt)
# print(corrected_prompt)

def detect_language(input_str):
    # 统计中文和英文字符的数量
    count_cn = count_en = 0
    for char in input_str:
        if '\u4e00' <= char <= '\u9fff':
            count_cn += 1
        elif char.isalpha():
            count_en += 1

    # 根据统计的字符数量判断主要语言
    if count_cn > count_en:
        return "cn"
    elif count_en > count_cn:
        return "en"
    else:
        return "unknow"




#定义Prompt文法
grammar = """
start: sentence
sentence: phrase ("," phrase)*
phrase: emphasis | weight | word | lora | embedding | schedule
emphasis: "(" sentence ")" -> emphasis
        | "[" sentence "]" -> weak_emphasis
weight: "(" word ":" NUMBER ")"
schedule: "[" word ":" word ":" NUMBER "]"
lora: "<" WORD ":" WORD (":" NUMBER)? (":" NUMBER)? ">"
embedding: "embedding" ":" WORD (":" NUMBER)? (":" NUMBER)?
word: WORD

NUMBER: /\s*-?\d+(\.\d+)?\s*/
WORD: /[^,:\(\)\[\]<>]+/
"""



@v_args(inline=True)  # Decorator to flatten the tree directly into the function arguments
class ChinesePromptTranslate(Transformer):
 
    def sentence(self, *args):
        return ", ".join(args)

    def phrase(self, *args):
        return "".join(args)

    def emphasis(self, *args):
        # Reconstruct the emphasis with translated content
        return "(" + "".join(args) + ")"

    def weak_emphasis(self, *args):
        print('weak_emphasis:',args)
        return "[" + "".join(args) + "]"
    
    def embedding(self,*args):
        print('prompt embedding',args[0])
        if len(args) == 1:
            # print('prompt embedding',str(args[0]))
            # 只传递了一个参数，意味着只有embedding名称没有数字
            embedding_name = str(args[0])
            return f"embedding:{embedding_name}"
        elif len(args) > 1:
            embedding_name,*numbers = args
           
            if len(numbers)==2:
                return f"embedding:{embedding_name}:{numbers[0]}:{numbers[1]}"
            elif len(numbers)==1:
                return f"embedding:{embedding_name}:{numbers[0]}"
            else:
                return f"embedding:{embedding_name}"

    def lora(self,*args):
        print('lora prompt',*args)
        if len(args) == 1:
            return f"<lora:{loar_name}>"
        elif len(args) > 1: 
            # print('lora', args)
            _,loar_name,*numbers = args
            loar_name = str(loar_name).strip()
            if len(numbers)==2:
                return f"<lora:{loar_name}:{numbers[0]}:{numbers[1]}>"
            elif len(numbers)==1:
                return f"<lora:{loar_name}:{numbers[0]}>"
            else:
                return f"<lora:{loar_name}>"
    
    def weight(self, word,number):
        translated_word = translate(str(word)).rstrip('.')
        return f"({translated_word}:{str(number).strip()})"
    
    def schedule(self,*args):
        print('prompt schedule',args)
        data = [str(arg).strip() for arg in args] 
        
        return f"[{':'.join(data)}]"

    def word(self, word):
        # Translate each word using the dictionary
        if detect_language(str(word)) == "cn":
            return  translate(str(word)).rstrip('.')
        else:
            return str(word).rstrip('.')
        
class ChinesePrompt:

    global _available
    available=_available
    
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
               "text": ("STRING",{"multiline": True,"default": "", "dynamicPrompts": False}),
               "generation": (["on","off"],{"default": "off"}), 
                             },
               
                "optional":{
                    "seed":("INT", {"default": 100, "min": 100, "max": 1000000}), 
                    
                },

                }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)

    FUNCTION = "run"

    CATEGORY = "♾️Mixlab/Prompt"
    OUTPUT_NODE = True
    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (True,)

    global text_pipe,zh_en_model,zh_en_tokenizer

    text_pipe= None
    zh_en_model=None
    zh_en_tokenizer=None 

    def run(self,text,seed,generation):
        

        seed=seed[0]
        generation=generation[0]

        # 进度条
        pbar = comfy.utils.ProgressBar(len(text)+1)
        texts = [correct_prompt_syntax(t) for t in text]

        global text_pipe,zh_en_model,zh_en_tokenizer
        if zh_en_model==None:
            zh_en_model = AutoModelForSeq2SeqLM.from_pretrained(zh_en_model_path).eval()
            zh_en_tokenizer = AutoTokenizer.from_pretrained(zh_en_model_path,padding=True, truncation=True)
        
        zh_en_model.to("cuda" if torch.cuda.is_available() else "cpu")
            # zh_en_tokenizer.to("cuda" if torch.cuda.is_available() else "cpu")
        
        text_pipe=pipeline('text-generation', model=text_generator_model_path,device="cuda" if torch.cuda.is_available() else "cpu")
        
        # text_pipe.model.to("cuda" if torch.cuda.is_available() else "cpu")

        prompt_result=[]

        # print('zh_en_model device',zh_en_model.device,text_pipe.model.device,torch.cuda.current_device() )
        en_texts=[]

        for t in texts:
            if t:
                # translated_text =  translated_word = translate(zh_en_tokenizer,zh_en_model,str(t))
                parser = Lark(grammar, start="start", parser="lalr", transformer=ChinesePromptTranslate())
                # print('t',t)
                result = parser.parse(t).children
                # print('en_result',result)        
                # en_text=translate(zh_en_tokenizer,zh_en_model,text_without_syntax)
                en_texts.append(result[0])

        zh_en_model.to('cpu')
        print("test en_text",en_texts)
        # en_text.to("cuda" if torch.cuda.is_available() else "cpu")

        pbar.update(1)
        for t in en_texts:
            if generation=='on':
                prompt =text_generate(text_pipe,t,seed)
                # 多条，还是单条
                lines = prompt.split("\n")
                longest_line = max(lines, key=len)
                # print(longest_line)
                prompt_result.append(longest_line)
            else:
                prompt_result.append(t)
            pbar.update(1)

        text_pipe.model.to('cpu')

        print('prompt_result',prompt_result,)
        # prompt_result = [','.join(correct_prompt_syntax(p)) for p in prompt_result]
        if len(prompt_result)==0:
            prompt_result=[""]
        return {
            "ui":{
                    "prompt": prompt_result
                },
            "result": (prompt_result,)}




class PromptGenerate:

    global _available
    available=_available
    
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "text": ("STRING",{"multiline": True,"default": "", "dynamicPrompts": False}),
                             },

                "optional":{
                    "multiple": (["off","on"],), 
                    "seed":("INT", {"default": 100, "min": 100, "max": 1000000}), 
                },

                }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)

    FUNCTION = "run"

    CATEGORY = "♾️Mixlab/Prompt"
    OUTPUT_NODE = True
    INPUT_IS_LIST = True
    OUTPUT_IS_LIST = (True,)

    global text_pipe

    text_pipe= None
    # 

    def run(self,text,multiple,seed):
        global text_pipe

        seed=seed[0]

        multiple=multiple[0]

        # 进度条
        pbar = comfy.utils.ProgressBar(len(text))
        
        text_pipe=pipeline('text-generation', model=text_generator_model_path,device="cuda" if torch.cuda.is_available() else "cpu")
        
        prompt_result=[]

        for t in text:
            prompt =text_generate(text_pipe,t,seed)
            prompt = prompt.split("\n")
            if multiple=='off':
                prompt = [max(prompt, key=len)]
            
            for p in prompt:
                prompt_result.append(p)
            pbar.update(1)

        text_pipe.model.to('cpu')

        return {
            "ui":{
                    "prompt": prompt_result
                },
            "result": (prompt_result,)}