diff --git a/__init__.py b/__init__.py index c76f9bd..6124cf1 100644 --- a/__init__.py +++ b/__init__.py @@ -12,18 +12,18 @@ from prompt_generator import PromptGenerator print("/_\ Loading Prompt Generator") -# Check prompt_generators folder under the models folder +# check prompt_generators folder under the models folder root = join(models_dir, "prompt_generators") if exists(root) is False: - print(f"/_\ {root} is created. Please add your prompt generators to {root} folder") mkdir(root) + print(f"/_\ {root} is created. Please add your prompt generators to {root} folder") prompts_file = join(base_path, "generated_prompts") if exists(prompts_file) is False: mkdir(prompts_file) -# Import PromptGenerator node to ComfyUI +# import PromptGenerator node to ComfyUI NODE_CLASS_MAPPINGS = { "Prompt Generator": PromptGenerator, @@ -34,3 +34,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +print("/_\ Loaded Successfully") +print("-" * 100) diff --git a/__pycache__/__init__.cpython-310.pyc b/__pycache__/__init__.cpython-310.pyc index 80e2ea1..94a085d 100644 Binary files a/__pycache__/__init__.cpython-310.pyc and b/__pycache__/__init__.cpython-310.pyc differ diff --git a/__pycache__/preprocess.cpython-310.pyc b/__pycache__/preprocess.cpython-310.pyc index 63d1872..4bf78a1 100644 Binary files a/__pycache__/preprocess.cpython-310.pyc and b/__pycache__/preprocess.cpython-310.pyc differ diff --git a/__pycache__/prompt_generator.cpython-310.pyc b/__pycache__/prompt_generator.cpython-310.pyc index e379518..830d1cf 100644 Binary files a/__pycache__/prompt_generator.cpython-310.pyc and b/__pycache__/prompt_generator.cpython-310.pyc differ diff --git a/generator/__init__.py b/generator/__init__.py index a6abef7..b9898ac 100644 --- a/generator/__init__.py +++ b/generator/__init__.py @@ -15,7 +15,7 @@ def check_package(package_name: str, install_name: str) -> None: print(f"/_\ Installing {package_name}") - command = "" + # check if portable version or manual if exists("python_embeded"): command = f".\\python_embeded\\python.exe -s -m pip install {install_name}" else: @@ -24,17 +24,20 @@ def check_package(package_name: str, install_name: str) -> None: process = run(command, shell=True, check=True, capture_output=True) +print(" Prompt Generator ComfyUI Node ".center(100, "-")) + # Check required packages print("/_\ Checking packages") check_package("transformers", "transformers") check_package("accelerate", "accelerate") +# triton package exists only in Linux if os_name == "Linux": check_package("triton", "triton") check_package("optimum", "optimum") -check_package("onnxruntime-gpu", "optimum[onnxruntime-gpu]") +check_package("onnxruntime", "optimum[onnxruntime-gpu]") """ # This package is for onnx models that run on CPU diff --git a/generator/__pycache__/__init__.cpython-310.pyc b/generator/__pycache__/__init__.cpython-310.pyc index dadf0bd..5694036 100644 Binary files a/generator/__pycache__/__init__.cpython-310.pyc and b/generator/__pycache__/__init__.cpython-310.pyc differ diff --git a/generator/__pycache__/generate.cpython-310.pyc b/generator/__pycache__/generate.cpython-310.pyc index ada5ffa..400dbfc 100644 Binary files a/generator/__pycache__/generate.cpython-310.pyc and b/generator/__pycache__/generate.cpython-310.pyc differ diff --git a/generator/__pycache__/model.cpython-310.pyc b/generator/__pycache__/model.cpython-310.pyc index a81fae7..879da89 100644 Binary files a/generator/__pycache__/model.cpython-310.pyc and b/generator/__pycache__/model.cpython-310.pyc differ diff --git a/generator/model.py b/generator/model.py index 14bb0bb..0725ad0 100644 --- a/generator/model.py +++ b/generator/model.py @@ -1,6 +1,8 @@ -from transformers import AutoModelForCausalLM, AutoTokenizer, Pipeline, pipeline +from transformers import AutoModelForCausalLM, AutoTokenizer, Pipeline + +from transformers import pipeline as tf_pipe +from optimum.pipelines import pipeline as opt_pipe -import optimum.pipelines from optimum.onnxruntime import ORTModelForCausalLM @@ -8,7 +10,7 @@ def get_default_pipeline(model_name: str) -> Pipeline: model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) - pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer) + pipe = tf_pipe(task="text-generation", model=model, tokenizer=tokenizer) return pipe @@ -17,7 +19,7 @@ def get_onnx_pipeline(model_name: str) -> Pipeline: model = ORTModelForCausalLM.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) - pipe = optimum.pipelines.pipeline( + pipe = opt_pipe( task="text-generation", model=model, tokenizer=tokenizer, accelerator="ort" ) @@ -28,7 +30,7 @@ def get_bettertransformer_pipeline(model_name: str) -> Pipeline: model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) - pipe = optimum.pipelines.pipeline( + pipe = opt_pipe( task="text-generation", model=model, tokenizer=tokenizer, diff --git a/preprocess.py b/preprocess.py index ff8bddd..5a0a475 100644 --- a/preprocess.py +++ b/preprocess.py @@ -1,54 +1,52 @@ from string import punctuation from re import sub, compile +from collections import OrderedDict -def get_unique_list(sequence : list) -> list: + +def get_unique_list(sequence: list) -> list: seen = set() return [x for x in sequence if not (x in seen or seen.add(x))] -def remove_exact_keywords(line : str) -> list[str]: + +def remove_exact_keywords(line: str) -> list[str]: char_blacklist = set(f"{punctuation}0123456789") # remove exact prompts prompts = get_unique_list(line.split(",")) - pure_prompts = [] + + pure_prompts = OrderedDict() # order matters + extracted_pure_prompts = {} # order isn't important # remove exact keyword for prompt in prompts: - can_add = True # extract the keyword keyword = "".join(c for c in prompt if c not in char_blacklist).lstrip() if keyword == "": continue - for pure_prompt in pure_prompts: - if prompt == pure_prompt: - can_add = False - break + if prompt in pure_prompts or keyword in extracted_pure_prompts: + continue - extracted_pure_prompt = "".join(c for c in pure_prompt if c not in char_blacklist).lstrip() - if keyword == extracted_pure_prompt: - can_add = False - break + pure_prompts[prompt] = True + extracted_pure_prompts[keyword] = True - if can_add: - pure_prompts.append(prompt) + return pure_prompts.keys() - return pure_prompts -def preprocess(line : str, preprocess_mode : str) -> str: - pattern = compile(r'(,\s){2,}') +def preprocess(line: str, preprocess_mode: str) -> str: + pattern = compile(r"(,\s){2,}") - temp_line = line.replace(u'\xa0', u' ') + temp_line = line.replace("\xa0", " ") temp_line = temp_line.replace("\n", ", ") temp_line = temp_line.replace("\t", " ") temp_line = temp_line.replace("|", ",") temp_line = temp_line.replace(" ", " ") - temp_line = sub(pattern, ', ', temp_line) + temp_line = sub(pattern, ", ", temp_line) if preprocess_mode == "exact_keyword": - temp_line = ','.join(remove_exact_keywords(temp_line)) + temp_line = ",".join(remove_exact_keywords(temp_line)) elif preprocess_mode == "exact_prompt": - temp_line = ','.join(get_unique_list(temp_line.split(","))) + temp_line = ",".join(get_unique_list(temp_line.split(","))) - return temp_line \ No newline at end of file + return temp_line diff --git a/prompt_generator.py b/prompt_generator.py index 76f9b41..063e2bf 100644 --- a/prompt_generator.py +++ b/prompt_generator.py @@ -16,7 +16,7 @@ class PromptGenerator: [ file for file in listdir(join(models_dir, "prompt_generators")) - if isdir(join(join(models_dir, "prompt_generators"), file)) + if isdir(join(models_dir, "prompt_generators", file)) ], ), "accelerate": (["enable", "disable"],),