diff --git a/README.md b/README.md index 7c3bdaa..8cff112 100644 --- a/README.md +++ b/README.md @@ -1,40 +1,55 @@ # ๐Ÿง  NeonLLama ComfyUI Extension -This custom ComfyUI node transforms a core **idea** into a richly detailed **positive prompt** using a local [Ollama](https://ollama.com) LLM. It also lets you specify **"avoid"** content, which is: - -- Used by the AI to **influence** the generated prompt by avoiding certain topics. -- Returned **unchanged** as the **negative prompt** for Stable Diffusion or similar models. +**NeonLLama** is a custom ComfyUI node that transforms one or more **idea lines** into vivid, richly detailed prompts using a local [Ollama](https://ollama.com) LLM. It also supports **avoid** content, which the model takes into account to steer generation โ€” and is returned as a **negative prompt**. --- ## ๐Ÿš€ Features -- ๐Ÿง  Generates a vivid, descriptive **positive prompt** from an idea. -- โ›” Lets you define what the AI should avoid mentioning (used during generation). -- ๐ŸŽฏ Returns your **avoid list directly as a negative prompt**, unmodified. -- ๐Ÿงฎ Token-aware generation using `clip-vit-base-patch32` tokenizer. -- ๐Ÿ” Retries until prompt fits within your token limits. -- โš™๏ธ Configurable parameters like token ranges, model, retry attempts, etc. +- ๐Ÿง  Generates a structured, richly visual **positive prompt** from a simple idea. +- โœ‚๏ธ Supports **multi-line idea inputs** โ€” each line becomes its own generation. +- โ›” Accepts an "avoid" list to influence prompt generation by avoiding unwanted terms. +- ๐ŸŽฏ Outputs the avoid list unmodified as the **negative prompt** (for Stable Diffusion, etc.). +- ๐Ÿงฎ Accurate token control using `laion/CLIP-ViT-bigG-14-laion2B-39B-b160k` tokenizer. +- ๐Ÿ” Retries intelligently until token limits are respected. +- โš™๏ธ Highly configurable: model choice, token ranges, max attempts, regen flag, and more. +- ๐Ÿ–จ๏ธ Returns all prompts in order and allows re-run on every use. --- ## ๐Ÿงฉ How It Works -1. You input an **idea** (e.g., `"cyberpunk alley in heavy rain"`). -2. You can add **avoid terms** (e.g., `"blur, soft lighting, extra limbs"`). -3. The LLM uses both to generate a positive prompt: - - The **idea** is expanded into a structured visual prompt. - - The **avoid terms** are used to **steer the generation away** from unwanted content. -4. The avoid list is also passed through untouched as the **negative prompt**. +1. You input one or more **ideas** (each on a new line). +2. You can optionally add **avoid terms** โ€” things to steer the model away from. +3. The model uses your idea(s) to generate descriptive prompts: + - **Positive Prompt**: Visually rich, structured, short-phrase based prompt. + - **Negative Prompt**: Your avoid terms, returned exactly as you wrote them. +4. Prompt generation is token-aware and dynamically revised if too long or too short. +5. Each idea line gets its own generated prompt; all are merged and returned. --- ## ๐Ÿ“ค Outputs -| Output | Type | Description | -|--------|------|-------------| -| `prompt` | `STRING` | The **positive prompt**, generated by the LLM. | -| `avoid` | `STRING` | The **negative prompt**, returned as provided. | +| Output | Type | Description | +|--------|--------|-----------------------------------------------------------------------------| +| `prompt` | STRING | Generated **positive prompt(s)**, joined with `BREAK` separators. | +| `negative` | STRING | The **avoid list**, passed through directly as the negative prompt. | +| `idea` | STRING | The original idea(s) input, for traceability or UI purposes. | + +--- + +## ๐Ÿ“ฅ Inputs / Configuration Fields + +| Name | Type | Description | +|--------------------|----------|-----------------------------------------------------------------------------| +| `model` | Dropdown | Select which locally running Ollama model to use. | +| `idea` | Text | The concept or image prompt idea. Up to 3 lines, one idea per line. | +| `negative` | Text | Words, elements, or themes to avoid. Used during gen + passed as negative. | +| `max_tokens` | Int | Max token limit (default: 75, max: 231). | +| `min_tokens` | Int | Minimum token requirement (default: 50). | +| `max_attempts` | Int | Number of tries to get an acceptable prompt per idea. | +| `regen_on_each_use`| Bool | Forces re-generation on each run, even if inputs didnโ€™t change. | --- @@ -42,41 +57,61 @@ This custom ComfyUI node transforms a core **idea** into a richly detailed **pos **Inputs:** -- `idea`: `haunted subway station with broken lights` -- `avoid`: `blood, gore, screaming` +```text +idea: + haunted subway station with broken lights + +avoid: + blood, gore, screaming +``` **Outputs:** -- `prompt`: *(LLM-generated)* - `"dark abandoned subway, flickering fluorescent lights, cracked tiled walls, shadowy corners, old train cars, graffiti-covered pillars, dim green glow, debris scattered floor"` +```text +prompt: + dark abandoned subway, flickering fluorescent lights, cracked tiled walls, shadowy corners, old train cars, graffiti-covered pillars, dim green glow, debris scattered floor -- `avoid`: *(Unchanged)* - `"blood, gore, screaming"` - -Use the `prompt` as your positive CLIP text, and `avoid` for negative conditioning. +negative: + blood, gore, screaming +``` --- -## โš™๏ธ Configuration Fields +## ๐Ÿ” Token-Aware Generation -| Name | Type | Description | -|------|------|-------------| -| `model` | Dropdown | Select the Ollama model to use. | -| `idea` | Multiline Text | The concept or image idea. | -| `avoid` | Multiline Text | Words/themes to avoid (used by LLM + passed to negative prompt). | -| `max_tokens` | Int | Maximum allowed tokens for the generated prompt. | -| `min_tokens` | Int | Minimum token target. | -| `max_attempts` | Int | Max retries to hit token range. | -| `regen_on_each_use` | Bool | Force prompt regeneration every time node runs. | +- Uses CLIP tokenizer (`laion/CLIP-ViT-bigG-14-laion2B-39B-b160k`) to measure tokens. +- Tries up to `max_attempts` times to reach your `min_tokens` and not exceed `max_tokens`. +- If prompt is too short, it adds detail; too long, it trims or rephrases. +- Avoid terms are included in a non-generative system message and tracked between retries. --- -## ๐Ÿ”„ Regeneration Behavior +## ๐Ÿง  Smart Prompt Structuring -If enabled, `regen_on_each_use` will force the node to re-generate the prompt every time it's executed, ensuring fresh output even if the inputs donโ€™t change. +The generation system instructs the LLM to: + +- Use **short phrases**, not sentences. +- **Avoid storytelling**, abstract ideas, and emotional language. +- Preserve the **core visual themes** of your original idea exactly. +- Avoid explanation or meta-commentary. +- Support basic structure with connectors like `with`, `under`, `surrounded by`. + +--- + +## ๐Ÿ› ๏ธ Requirements + +- [Ollama](https://ollama.com) installed and running locally (on `localhost:11434`). +- ComfyUI environment. +- Python dependencies: `requests`, `tokenizers`. --- ## ๐Ÿ“„ License MIT License + +--- + +## โค๏ธ Credits + +Built by Neon Lightning โšก \ No newline at end of file diff --git a/__init__.py b/__init__.py index 158cc7d..2ab6083 100644 --- a/__init__.py +++ b/__init__.py @@ -21,7 +21,7 @@ def fetch_ollama_models(): OLLAMA_MODELS = fetch_ollama_models() -tokenizer = Tokenizer.from_pretrained("openai/clip-vit-base-patch32") +tokenizer = Tokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") def estimate_tokens(text): return tokenizer.encode(text).ids @@ -32,16 +32,16 @@ class OllamaPromptFromIdea: return { "required": { "model": (OLLAMA_MODELS, {"tooltip": "Select the Ollama model to generate prompts with."}), - "idea": ("STRING", {"multiline": True, "default": "futuristic cyberpunk city", "tooltip": "Enter the core concept or theme for your prompt to ollama."}), - "avoid": ("STRING", {"multiline": True, "default": "", "tooltip": "Words or themes to exclude from the prompt.(non ollama prompting)"}), - "max_tokens": ("INT", {"default": 77, "min": 10, "max": 231, "tooltip": "Maximum token length for the generated prompt."}), - "min_tokens": ("INT", {"default": 60, "min": 10, "max": 230, "tooltip": "Minimum token length for the generated prompt."}), + "idea": ("STRING", {"multiline": True, "default": "futuristic cyberpunk city", "tooltip": "Enter the core concept or theme for your prompt to ollama\nYou can have seperated ideas if you have a hard return.\nOnly use up to 3 lines though. to a maximum. of 231 tokens."}), + "negative": ("STRING", {"multiline": True, "default": "", "tooltip": "Words or themes to exclude from the prompt.(non ollama prompting)"}), + "max_tokens": ("INT", {"default": 75, "min": 10, "max": 231, "tooltip": "Maximum token length for the generated prompt."}), + "min_tokens": ("INT", {"default": 50, "min": 10, "max": 230, "tooltip": "Minimum token length for the generated prompt."}), "max_attempts": ("INT", {"default": 30, "min": 1, "max": 200, "tooltip": "Number of attempts to generate a prompt fitting token limits."}), "regen_on_each_use": ("BOOLEAN", {"default": True, "tooltip": "Force regeneration on each node execution."}), } } - RETURN_TYPES = ("STRING", "STRING",) - RETURN_NAMES = ("prompt", "avoid",) + RETURN_TYPES = ("STRING", "STRING", "STRING",) + RETURN_NAMES = ("prompt", "negative", "idea") FUNCTION = "generate_prompt" CATEGORY = "Ollama" @@ -51,112 +51,135 @@ class OllamaPromptFromIdea: return float("NaN") return None - def generate_prompt(self, model, idea, avoid, max_tokens, min_tokens, max_attempts, regen_on_each_use): - if not avoid: - avoid = "" + def generate_prompt(self, model, idea, negative, max_tokens, min_tokens, max_attempts, regen_on_each_use): + if not negative: + negative = "" token_min = min(min_tokens, max_tokens) token_expand_threshold = int(token_min * 0.75) - last_output = None - avoid_clause = f"\nABSOLUTELY avoid mentioning: {avoid.strip()}" if avoid.strip() else "" - for attempt in range(1, max_attempts + 1): - try: - if last_output is None: - system_prompt = ( - f"Convert the following idea into a richly descriptive, visually detailed image prompt for Stable Diffusion XL. " - f"Use short but expressive phrases, and allow natural connectors like 'with', 'and', or 'under'. " - f"Focus on concrete, vivid visual elements โ€“ not abstract concepts. " - f"Use multi-word descriptions where appropriate. " - f"Do not include full sentences, storytelling, or subjective opinions. " - f"Use only short descriptions. and don't describe feeling" - f"Avoid overly generic or disconnected terms. " - f"Target between {token_min} and {max_tokens} tokens. " - f"{avoid_clause}" - f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {idea} DO NOT CHANGE THE IDEA." - f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." - f"\nIdea: {idea}\nPrompt:" - ) - else: - token_count = len(estimate_tokens(last_output)) - if token_count > max_tokens: - system_prompt = ( - f"The following prompt is too long (over {max_tokens} tokens). " - f"Revise it to be shorter but keep visual richness and specificity. " - f"Use compact phrases or brief expressions with light structure. " - f"Avoid long sentences or reinterpreting the concept. " - f"Use only short descriptions. and don't describe feeling" - f"{avoid_clause}" - f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {idea} DO NOT CHANGE THE IDEA." - f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." - f"\nPrevious prompt: {last_output}\nShorter prompt:" + + idea_list = [i.strip() for i in idea.strip().split("\n") if i.strip()] + generated_prompts = [] + + for idx, sub_idea in enumerate(idea_list): + print(f"\n๐Ÿง  Generating prompt for idea {idx + 1}: '{sub_idea}'") + last_output = None + used_phrases = [] + if negative.strip(): + used_phrases.append(negative.strip()) + + for attempt in range(1, max_attempts + 1): + try: + avoid_text = " | ".join(used_phrases) + avoid_clause = "" + if avoid_text.strip() and (negative.strip() or idx > 0): + avoid_clause = ( + f"\nABSOLUTELY avoid using or repeating any of the following phrases or content but keep them in mind: {avoid_text}" ) - elif token_count < token_expand_threshold: + if last_output is None: system_prompt = ( - f"The following prompt is too short (under {token_expand_threshold} tokens). " - f"Expand it by adding specific, vivid imagery using short but rich phrases. " - f"Include unique textures, lighting effects, environments, and visual motifs. " - f"Light structure is allowed: use connectors like 'with', 'under', 'surrounded by', etc. " - f"Do not repeat phrases or rearrange words โ€“ add new coherent, visual material. " - f"Use only short descriptions. and don't describe feeling" - f"Avoid full sentences or storylines. " + f"Convert the following idea into a richly descriptive, visually detailed image prompt for Stable Diffusion XL. " + f"Use short phrases, and allow natural connectors like 'with', 'and', or 'under'. " + f"Focus on concrete, vivid visual elements โ€“ not abstract concepts. " + f"Use multi-word descriptions only where needed. " + f"Do not include full sentences, storytelling, or subjective opinions. " + f"Use only short descriptions. and don't describe feeling. " + f"Target between {token_min} and {max_tokens} tokens. " f"{avoid_clause}" - f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {idea} DO NOT CHANGE THE IDEA." + f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {sub_idea} DO NOT CHANGE THE IDEA." f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." - f"\nPrevious prompt: {last_output}\nExpanded prompt:" + f"\nIdea: {sub_idea}\nPrompt:" ) else: - system_prompt = ( - f"Revise the following prompt to improve clarity and vividness, while keeping all original ideas intact. " - f"You may slightly structure the phrases for better flow. " - f"Do not add new concepts or remove core elements. " - f"Use only short descriptions. and don't describe feeling" - f"Keep it between {token_min}โ€“{max_tokens} tokens. " - f"{avoid_clause}" - f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {idea} DO NOT CHANGE THE IDEA." - f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." - f"\nPrevious prompt: {last_output}\nRevised prompt:" - ) - seed = random.randint(0, 99999999) - response = requests.post( - "http://localhost:11434/api/generate", - json={ - "model": model, - "prompt": system_prompt, - "stream": False, - "options": { - "seed": seed, - "temperature": 0.7 - } - }, - timeout=120, - ) - response.raise_for_status() - raw_result = response.json().get("response", "").strip() - token_count = len(estimate_tokens(raw_result)) - print(f"Attempt {attempt} Ollama result: {raw_result}") - print(f"โ†’ Token count: {token_count} (target: {token_min}โ€“{max_tokens})") - if last_output is not None and raw_result.strip() == last_output.strip(): - print("โš ๏ธ Prompt identical to last attempt. Restarting generation from scratch...\n") - last_output = None + token_count = len(estimate_tokens(last_output)) + if token_count > max_tokens: + system_prompt = ( + f"The following prompt is too long (over {max_tokens} tokens). " + f"Revise it to be shorter but keep visual richness and specificity. " + f"Use compact phrases or brief expressions with light structure. " + f"Avoid long sentences or reinterpreting the concept. " + f"Use only short descriptions. and don't describe feeling. " + f"{avoid_clause}" + f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {sub_idea} DO NOT CHANGE THE IDEA." + f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." + f"\nPrevious prompt: {last_output}\nShorter prompt:" + ) + elif token_count < token_expand_threshold: + system_prompt = ( + f"The following prompt is too short (under {token_expand_threshold} tokens). " + f"Expand it by adding specific, vivid imagery using short but rich phrases. " + f"Include unique textures, lighting effects, environments, and visual motifs. " + f"Light structure is allowed: use connectors like 'with', 'under', 'surrounded by', etc. " + f"Do not repeat phrases or rearrange words โ€“ add new coherent, visual material. " + f"Use only short descriptions. and don't describe feeling. " + f"Avoid full sentences or storylines. " + f"{avoid_clause}" + f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {sub_idea} DO NOT CHANGE THE IDEA." + f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." + f"\nPrevious prompt: {last_output}\nExpanded prompt:" + ) + else: + system_prompt = ( + f"Revise the following prompt to improve clarity and vividness, while keeping all original ideas intact. " + f"You may slightly structure the phrases for better flow. " + f"Do not add new concepts or remove core elements. " + f"Use only short descriptions. and don't describe feeling. " + f"Keep it between {token_min}โ€“{max_tokens} tokens. " + f"{avoid_clause}" + f"Reminder: You MUST preserve all core themes of the original idea. The original idea is: {sub_idea} DO NOT CHANGE THE IDEA." + f"you MUST NOT ever talk about your thought process or explain how you generated the prompt." + f"\nPrevious prompt: {last_output}\nRevised prompt:" + ) + seed = random.randint(0, 99999999) + response = requests.post( + "http://localhost:11434/api/generate", + json={ + "model": model, + "prompt": system_prompt, + "stream": False, + "options": { + "seed": seed, + "temperature": 0.7 + } + }, + timeout=120, + ) + response.raise_for_status() + raw_result = response.json().get("response", "").strip() + token_count = len(estimate_tokens(raw_result)) + print(f"Idea: {idx + 1} Attempt: {attempt}/{max_attempts} Ollama result: {raw_result}") + print(f"โ†’ Token count: {token_count} (target: {token_min}โ€“{max_tokens})") + + if last_output is not None and raw_result.strip() == last_output.strip(): + print("โš ๏ธ Prompt identical to last attempt. Restarting generation from scratch...\n") + last_output = None + seed = random.randint(0, 99999999) + time.sleep(0.5) + continue + + if token_min <= token_count <= max_tokens: + used_phrases.append(raw_result) + print("โœ”๏ธ Prompt accepted.") + generated_prompts.append(raw_result) + break + last_output = raw_result + reason = "too long" if token_count > max_tokens else "too short" + print(f"โš ๏ธ Prompt {reason}. Retrying...\n") + time.sleep(0.5) + except requests.exceptions.Timeout: + print(f"โš ๏ธ Attempt {attempt}/{max_attempts} timed out. Retrying...\n") time.sleep(0.5) continue - if token_min <= token_count <= max_tokens: - print("โœ”๏ธ Prompt accepted.") - return (raw_result, avoid) - last_output = raw_result - reason = "too long" if token_count > max_tokens else "too short" - print(f"โš ๏ธ Prompt {reason}. Retrying...\n") - time.sleep(0.5) - except requests.exceptions.Timeout: - print(f"โš ๏ธ Attempt {attempt} timed out. Retrying...\n") - time.sleep(0.5) - continue - except Exception as e: - error_msg = f"[Ollama Error] {str(e)}" - print(error_msg) - return (error_msg, "") - print("โŒ Max attempts reached. Returning fallback idea.") - return (idea, avoid) + except Exception as e: + error_msg = f"[Ollama Error] {str(e)}" + print(error_msg) + return (error_msg, "") + else: + print(f"โŒ Max attempts for idea '{sub_idea}' reached. Using original as fallback.") + generated_prompts.append(sub_idea) + outputend="\nBREAK\n".join(generated_prompts) + print(f"output of:{outputend}") + return (" BREAK ".join(generated_prompts), negative, idea) def ui(self, inputs, outputs): prompt_str = outputs[0] if isinstance(outputs, (list, tuple)) and outputs else ""