commit fc08051f240ecf8206deeb1f285e0242174452a7 Author: NeuralSamurAI <164533106+NeuralSamurAI@users.noreply.github.com> Date: Wed Aug 7 15:36:24 2024 -0400 Add files via upload diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2ccfb1b --- /dev/null +++ b/__init__.py @@ -0,0 +1,14 @@ +from .prompt_json import PromptJSON +from .token_counter import TokenCounter + +NODE_CLASS_MAPPINGS = { + "PromptJSON": PromptJSON, + "TokenCounter": TokenCounter +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PromptJSON": "Prompt JSON", + "TokenCounter": "Token Counter" +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..59c2abb --- /dev/null +++ b/readme.md @@ -0,0 +1,217 @@ +# PromptJSON Node for ComfyUI + +## Description +PromptJSON is a custom node for ComfyUI that structures natural language prompts and generates prompts for external LLM nodes in image generation workflows. It aids in creating consistent, schema-based image descriptions. + +## Features +- Structures natural language prompts according to a custom schema +- Generates prompts for external LLM nodes with multiple strategies +- Supports negative prompts +- Adjustable complexity for output detail +- Custom color scheme input +- Custom schema support for fine-tuning prompt structure +- Multiple prompt generation strategies +- Token counting for T5 and CLIP models +- Batch processing support +- Debug mode for detailed information +- Error handling for robust operation + +## Installation +1. Navigate to your ComfyUI custom nodes directory. +2. Clone this repository or copy the `prompt_json.py`, `token_counter.py`, and `__init__.py` files into a new directory named `PromptJSON`. +3. Install required dependencies: + ``` + pip install nltk + ``` +4. Restart ComfyUI or refresh your workflow. + +## Usage +1. Find the "Prompt JSON" node in the "prompt_converters" category in ComfyUI. +2. Connect the following inputs: + - `prompt`: Your main prompt text + - `negative_prompt`: Elements to avoid in the image + - `complexity`: A float value between 0.1 and 1.0 to adjust output detail + - `max_tokens`: Maximum number of tokens for the output JSON + - `model`: The model to use for tokenization + - `llm_prompt_type`: Choose between "One Shot Learning", "System Prompt", or "Combined" + - `color_scheme` (optional): Comma-separated list of colors + - `seed` (optional): Seed for random number generation (-1 for random seed) + - `custom_schema` (optional): JSON string for custom schema + - `debug_mode` (optional): Boolean to enable debug information in output + - `batch_size` (optional): Number of outputs to generate in a batch + +3. Use the outputs: + - `llm_prompt`: The prompt to be sent to the external LLM node + - `structured_prompt`: A JSON representation of the input prompt, structured according to the schema + - `custom_schema`: The schema used to structure the prompt and guide the LLM + +## LLM Prompt Types +- **One Shot Learning**: Provides an example of a structured response based on the input prompt. +- **System Prompt**: Gives instructions to the LLM on how to generate structured descriptions. +- **Combined**: Merges the System Prompt and One Shot Learning example for potentially increased accuracy. + +## Workflow +1. The PromptJSON node takes the user's input and structures it according to the custom schema. +2. It generates an LLM prompt based on the chosen prompt type, incorporating the structured prompt and schema. +3. The `llm_prompt` output is sent to an external LLM node. +4. The external LLM node generates the final, detailed JSON description based on this prompt. + +## Custom Schema +The custom schema feature allows you to define the structure of the output JSON and influence the generated LLM prompt. + +### Example: Custom Schema for Image Generation +``` +{ + "scene": { + "time_of_day": "string", + "weather": "string", + "location": "string" + }, + "subjects": [ + { + "type": "string", + "description": "string", + "position": "string" + } + ], + "style": { + "artistic_movement": "string", + "color_palette": ["string"], + "mood": "string" + }, + "camera": { + "angle": "string", + "shot_type": "string" + } +} +``` + +This schema will guide the structuring of the input prompt and the generation of the LLM prompt. + +## Example Output +Input: +- prompt: "A serene lake at sunset with a lone fisherman in a small boat" +- negative_prompt: "urban, city, crowds" +- complexity: 0.7 +- max_tokens: 2048 +- color_scheme: "orange, purple, blue" +- custom_schema: (The schema shown above) +- llm_prompt_type: "Combined" + +Output: +1. LLM Prompt: +``` +You are an AI assistant specialized in generating detailed image descriptions based on the following schema: +{ + "scene": { + "time_of_day": "string", + "weather": "string", + "location": "string" + }, + "subjects": [ + { + "type": "string", + "description": "string", + "position": "string" + } + ], + "style": { + "artistic_movement": "string", + "color_palette": ["string"], + "mood": "string" + }, + "camera": { + "angle": "string", + "shot_type": "string" + } +} + +Here's an example of how to structure your response: + +{ + "scene": { + "time_of_day": "sunset", + "weather": "clear", + "location": "serene lake" + }, + "subjects": [ + { + "type": "character", + "description": "lone fisherman", + "position": "in a small boat" + }, + { + "type": "object", + "description": "small boat", + "position": "on the lake" + } + ], + "style": { + "artistic_movement": "Impressionism", + "color_palette": ["orange", "purple", "blue"], + "mood": "tranquil" + }, + "camera": { + "angle": "wide shot", + "shot_type": "landscape" + } +} + +Please generate a similar structured description for the given prompt, but with more detail and creativity. + +Generate a detailed image description based on the following prompt, adhering to the provided schema: +{ + "scene": { + "time_of_day": "sunset", + "weather": "clear", + "location": "serene lake" + }, + "subjects": [ + { + "type": "character", + "description": "lone fisherman", + "position": "in a small boat" + }, + { + "type": "object", + "description": "small boat", + "position": "on the lake" + } + ], + "style": { + "artistic_movement": "Impressionism", + "color_palette": ["orange", "purple", "blue"], + "mood": "tranquil" + }, + "camera": { + "angle": "wide shot", + "shot_type": "landscape" + } +} +``` + +2. Structured Prompt: (The JSON representation of the input prompt, as shown in the LLM Prompt) + +3. Custom Schema: (The schema used, as shown at the beginning of the LLM Prompt) + +## Notes +- The `structured_prompt` output is an intermediate representation, not the final detailed JSON. +- The `llm_prompt` is designed to guide an external LLM in generating the final, detailed structured image description. +- Experiment with different LLM prompt types to find the most effective one for your specific LLM and use case. +- The custom schema defines the structure for both the intermediate prompt and the final LLM output. +- Adjust the complexity parameter to control the level of detail in the generated prompts. + +## Token Counter Usage +The Token Counter node is available in the "utils" category: + +1. Connect the following inputs: + - `text`: The text to count tokens for + - `model`: The model to use for tokenization + - `use_clip` (optional): Boolean to enable CLIP token counting + +2. Use the outputs: + - `t5_token_count`: Number of tokens using T5 tokenizer + - `clip_token_count`: Number of tokens using CLIP tokenizer (if enabled) + +## Support +For issues, feature requests, or contributions, please open an issue or pull request in the GitHub repository. \ No newline at end of file diff --git a/token_counter.py b/token_counter.py new file mode 100644 index 0000000..aae7216 --- /dev/null +++ b/token_counter.py @@ -0,0 +1,49 @@ +from transformers import T5Tokenizer, CLIPTokenizer + +class TokenCounter: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING", {"multiline": True}), + "model": ("MODEL",), + }, + "optional": { + "use_clip": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("t5_token_count", "clip_token_count") + FUNCTION = "count_tokens" + + CATEGORY = "utils" + + def count_tokens(self, text, model, use_clip=False): + if not isinstance(text, str): + raise ValueError("Input text must be a string") + + t5_tokenizer = self.get_t5_tokenizer(model) + t5_tokens = t5_tokenizer.encode(text) + t5_count = len(t5_tokens) + + if use_clip: + clip_tokenizer = self.get_clip_tokenizer(model) + clip_tokens = clip_tokenizer.encode(text) + clip_count = len(clip_tokens) + else: + clip_count = None + + return t5_count, clip_count + + def get_t5_tokenizer(self, model): + try: + return model.model.t5_tokenizer + except AttributeError: + return T5Tokenizer.from_pretrained("t5-xxl", local_files_only=True) + + def get_clip_tokenizer(self, model): + try: + return model.model.clip_tokenizer + except AttributeError: + return CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", local_files_only=True)