diff --git a/README.md b/README.md index c280adb..d0e9c0f 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ Custom prompt generator node for ComfyUI # Table Of Contents +- [prompt-generator-comfyui](#prompt-generator-comfyui) +- [Table Of Contents](#table-of-contents) - [Setup](#setup) - [Features](#features) - [Example Workflow](#example-workflow) @@ -11,15 +13,13 @@ Custom prompt generator node for ComfyUI - [Example Outputs](#example-outputs) # Setup -- Run ```pip install happytransformer``` command in the environment that you are launching ComfyUI with -- Copy ```prompt_generator.py``` file to ```custom_nodes``` folder in ComfyUI -- Create ```prompt_generators``` folder under ```models``` folder in ComfyUI +- Clone the repository with ```https://github.com/alpertunga-bile/prompt-generator-comfyui.git``` command under ```custom_nodes``` folder. - Put your generator under ```prompt_generators``` folder. You can create your prompt generator with [this repository](https://github.com/alpertunga-bile/prompt-markdown-parser). You have to put generator as folder. Do not just put ```pytorch_model.bin``` file for example. - Run the ComfyUI - Open the ```hires.fixWithPromptGenerator.json``` workflow # Features -- Print generated text to terminal and log the node's state in ```generated_prompts.txt``` file +- Print generated text to terminal and log the node's state under ```generated_prompts``` folder with date as filename. # Example Workflow ![example_workflow](https://github.com/alpertunga-bile/prompt-generator-comfyui/assets/76731692/f50652a9-8751-41f3-81cf-d4cb61dd8a34) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8c9d8b4 --- /dev/null +++ b/__init__.py @@ -0,0 +1,54 @@ +from sys import path +from os.path import dirname, exists, join +from subprocess import run +from os import remove, mkdir + +path.append(dirname(__file__)) + +from prompt_generator import PromptGenerator + +print("/_\ Loading Prompt Generator") + +# Check prompt_generators folder under models folder + +root = join("models", "prompt_generators") +if exists(root) is False: + print(f"/_\ {root} is created. Please add your prompt generators to {root} folder") + mkdir(root) + +if exists("generated_prompts") is False: + mkdir("generated_prompts") + +# Check happytranformer package + +temp_requirements_file = "temp_requirements.txt" + +process = run(f"pip freeze > {temp_requirements_file}", shell=True, check=True, capture_output=True) +need_to_install = True +packages = set() + +with open(temp_requirements_file, "r") as file: + packages = set(file.readlines()) + +for package in packages: + if "happytransformer" in package: + need_to_install = False + break + +remove(temp_requirements_file) + +if need_to_install: + print("/_\ Installing happytransformer") + process = run("pip install happytransformer", shell=True, check=True, capture_output=True) + +# Import PromptGenerator node to ComfyUI + +NODE_CLASS_MAPPINGS = { + "Prompt Generator": PromptGenerator +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Prompt Generator": "Prompt Generator" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/prompt_generator.py b/prompt_generator.py index 29ec9b5..4f8f8d5 100644 --- a/prompt_generator.py +++ b/prompt_generator.py @@ -2,12 +2,6 @@ from os import listdir, mkdir from os.path import join, isdir, exists class PromptGenerator: - def __init__(self) -> None: - root = join("models", "prompt_generators") - if exists(root) is False: - print(f"{root} is created. Please add your prompt generators to {root} folder") - mkdir(root) - @classmethod def INPUT_TYPES(s): return { @@ -187,10 +181,11 @@ class PromptGenerator: def generate(self, clip, model_type, model_name, seed, min_length, max_length, do_sample, early_stopping, num_beams, temperature, top_k, top_p, no_repeat_ngram_size, self_recursive, recursive_level, preprocess_mode): from happytransformer import HappyGeneration, GENSettings + from datetime import date root = join("models", "prompt_generators") real_path = join(root, model_name) - prompt_log_filename = "generated_prompts.txt" + prompt_log_filename = join("generated_prompts", str(date.today())) generated_text = "" if exists(prompt_log_filename) is False: @@ -227,13 +222,4 @@ class PromptGenerator: tokens = clip.tokenize(generated_text) cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) - return ([[cond, {"pooled_output": pooled}]], ) - -NODE_CLASS_MAPPINGS = { - "Prompt Generator": PromptGenerator -} - -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = { - "Prompt Generator": "Prompt Generator" -} + return ([[cond, {"pooled_output": pooled}]], ) \ No newline at end of file