import random import subprocess import requests import time import requests.exceptions from pathlib import Path from comfy.utils import ProgressBar import threading import traceback LMSTUDIO_AVAILABLE = False lms = None try: import lmstudio as lms from lmstudio.sync_api import Client LMSTUDIO_AVAILABLE = True except ImportError: print("⚠️ LM Studio SDK not installed. LM Studio functionality disabled.") ALL_AVAILABLE_MODELS = ["Loading models..."] tokenizer = None MODELS_LOADED = False def initialize_in_background(): global ALL_AVAILABLE_MODELS, tokenizer, MODELS_LOADED models = [] try: res = requests.get("http://localhost:11434/api/tags", timeout=1.5) if res.status_code == 200: ollama_models = [m["model"] for m in res.json().get("models", [])] models.extend([f"ollama:{m}" for m in ollama_models]) print(f"✅ Found {len(ollama_models)} Ollama models") except Exception as e: print(f"⚠️ Ollama model fetch skipped: {str(e)}") if LMSTUDIO_AVAILABLE and lms is not None: try: client = Client() result = [] def fetch_models(): try: nonlocal result downloaded_models = client.list_downloaded_models() or [] result = [ m.model_key for m in downloaded_models if isinstance(m, lms.DownloadedLlm) ] except Exception as e: print(f"LM Studio fetch error: {str(e)}") result = [] t = threading.Thread(target=fetch_models) t.daemon = True t.start() t.join(timeout=2.0) if t.is_alive(): print("⚠️ LM Studio model fetch timed out after 2 seconds") else: models.extend([f"lmstudio:{key}" for key in result]) print(f"✅ Found {len(result)} LM Studio models") except Exception as e: print(f"⚠️ LM Studio model fetch failed: {str(e)}") traceback.print_exc() if models: ALL_AVAILABLE_MODELS = models else: ALL_AVAILABLE_MODELS = ["No models available - check Ollama/LM Studio"] try: from tokenizers import Tokenizer tokenizer = Tokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") print("✅ Tokenizer loaded successfully") except Exception as e: print(f"⚠️ Tokenizer load failed: {str(e)}") tokenizer = None MODELS_LOADED = True threading.Thread(target=initialize_in_background, daemon=True).start() class OllamaPromptFromIdea: @classmethod def IS_CHANGED(cls, **kwargs): return float("NaN") @classmethod def INPUT_TYPES(cls): llm_seed = random.randint(0, 1000000000) llm_temp = random.uniform(0.1, 1.0) return { "required": { "model": (ALL_AVAILABLE_MODELS, {"tooltip": "Select the LLM model (Ollama or LM Studio) to generate prompts with."}), "idea": ("STRING", {"multiline": True, "default": "futuristic cyberpunk city", "tooltip": "Enter the core concept or theme for your prompt.\nYou can have separated 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 (used by Stable Diffusion, not LLM).", "dynamicPrompts": False}), "max_tokens": ("INT", {"default": 75, "min": 10, "max": 1024, "tooltip": "Maximum token length for the generated prompt."}), "min_tokens": ("INT", {"default": 50, "min": 10, "max": 1024, "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 (doesn't matter if just_use_idea is on)."}), "just_use_idea": ("BOOLEAN", {"default": False, "tooltip": "Skip Generating and just use idea as prompt."}), "exclude_comma": ("BOOLEAN", {"default": False, "tooltip": "Disables commas and sentence removal suggesting."}), "randomize_seed": ("BOOLEAN", {"default": True, "tooltip": "Use a random seed on each generation."}), "llm_seed": ("INT", {"default": llm_seed, "min": 0, "max": 999999999, "tooltip": "Fixed seed (only used if randomize_seed is off)."}), "randomize_temp": ("BOOLEAN", {"default": True, "tooltip": "Use a random temperature on each generation."}), "llm_temp": ("FLOAT", {"default": llm_temp, "min": 0.1, "max": 1, "tooltip": "Fixed temperature (only used if randomize_temp is off)."}), "keepllm": ("BOOLEAN", {"default": False, "tooltip": "Keep LLM model in memory."}), } } RETURN_TYPES = ("STRING", "STRING", "STRING",) RETURN_NAMES = ("prompt", "negative", "idea") FUNCTION = "generate_prompt" CATEGORY = "LLM Prompts" def generate_prompt(self, model, idea, negative, max_tokens, min_tokens, max_attempts, regen_on_each_use, just_use_idea, exclude_comma, randomize_seed, llm_seed, randomize_temp, llm_temp, keepllm): PREDICTION_TIMEOUT = 120 if just_use_idea: print("[LLM Prompt Node] 'Just Use Idea' is enabled. Skipping LLM generation.") return (idea, negative, idea) if not negative: negative = "" is_ollama_model = model.startswith("ollama:") is_lmstudio_model = model.startswith("lmstudio:") if not is_ollama_model and not is_lmstudio_model: print(f"[LLM Prompt Node] Error: Invalid model selection prefix for '{model}'. Falling back to idea.") return (idea, negative, idea) actual_model_name = model.split(":", 1)[1] if ":" in model else model print(f"[LLM Prompt Node] Using model: {actual_model_name} from {'Ollama' if is_ollama_model else 'LM Studio'}") idea_list = [i.strip() for i in idea.strip().split("\n") if i.strip()] total_attempts = len(idea_list) * max_attempts pbar = ProgressBar(total_attempts) if not idea_list: print("[LLM Prompt Node] No valid ideas provided. Returning empty prompt.") return ("", negative, "") generated_prompts = [] prompt_log_file = Path("llm_generated_prompt_log.txt") if not regen_on_each_use: try: if prompt_log_file.exists(): with open(prompt_log_file, "r", encoding="utf-8") as f: file_prompt = f.read().strip() print("[LLM Prompt Node] Loaded prompt from file instead of regenerating.") return (file_prompt, negative, idea) else: print("[LLM Prompt Node] Prompt log file not found. Generating new prompt.") except Exception as e: print(f"[LLM Prompt Node] Failed to load cached prompt file: {e}. Generating new prompt.") lm_studio_llm_instance = None if is_lmstudio_model: try: lm_studio_client = lms.Client() lm_studio_llm_instance = lm_studio_client.llm.model(actual_model_name) except Exception as e: print(f"[LLM Prompt Node] Error initializing LM Studio client or model: {e}") return (idea, negative, idea) for idx, sub_idea in enumerate(idea_list): print(f"\n🧠 Generating prompt for idea {idx + 1}: '{sub_idea}'") last_output = None used_phrases = [] timeout_number = 0 if negative.strip(): used_phrases.append(negative.strip()) for attempt in range(1, max_attempts + 1): pbar.update_absolute((idx * max_attempts) + (attempt - 1)) try: avoid_text = " | ".join(used_phrases) avoid_clause = "" if avoid_text.strip(): avoid_clause = f"\nABSOLUTELY DO NOT use or repeat any of the following phrases or content: {avoid_text}." if avoid_text.strip() else "" if exclude_comma: system_message_content = ( f"You are a specialized prompt generator for Stable Diffusion. " f"Your task is to convert a raw idea into a single-line, visually dense, and concrete image prompt. " f"Use only short, descriptive fragments – no emotions, no abstract terms, no opinions. " f"Focus strictly on visible elements: subject appearance, setting, lighting, objects, materials, and structure. " f"Use connectors like 'with', 'under', 'surrounded by', but avoid excessive chaining. " f"Sort ideas by visual importance, from main subject to secondary elements. " f"The entire prompt must be between {min_tokens} and {max_tokens} tokens. " f"{avoid_clause}" f"DO NOT explain your reasoning or describe what you are doing. DO NOT reinterpret the original idea. Preserve it: {sub_idea}" f"Never use storytelling, feelings, or narrative context." f"Your output must be suitable for direct input into an AI image generation model." f"Give a new idea from any previous" ) else: system_message_content = ( f"You are a specialized prompt generator for Stable Diffusion. " f"Your task is to convert a raw idea into a single-line, visually dense, and concrete image prompt. " f"Use only short, descriptive fragments – no full sentences, no emotions, no abstract terms, no opinions. " f"Focus strictly on visible elements: subject appearance, setting, lighting, objects, materials, and structure. " f"Use connectors like 'with', 'under', 'surrounded by', but avoid excessive chaining. " f"Sort ideas by visual importance, from main subject to secondary elements. " f"Separate all ideas with commas. In order of importance." f"The entire prompt must be between {min_tokens} and {max_tokens} tokens" f"{avoid_clause}" f"DO NOT explain your reasoning or describe what you are doing. DO NOT reinterpret the original idea. Preserve it: {sub_idea}" f"Never use storytelling, feelings, or narrative context." f"Your output must be suitable for direct input into an AI image generation model." f"Give a new idea from any previous" ) user_message_content = f"Idea: {sub_idea}\nPrompt:" if last_output is not None: token_count_last_output = estimate_tokens(last_output) if token_count_last_output > max_tokens: user_message_content = ( f"DO NOT explain your reasoning or describe what you are doing." f"The following prompt is too long it needs to be slightly shorter." f"Shorten it without removing detail. Compress phrases, remove redundancy. " f"{avoid_clause}" f"Never use storytelling, feelings, or narrative context." f"Your output must be suitable for direct input into an AI image generation model." f"Give a new idea from the previous" f"{last_output}" ) elif token_count_last_output < min_tokens: user_message_content = ( f"DO NOT explain your reasoning or describe what you are doing." f"The following prompt is too short it needs to be slightly longer.. " f"Expand it with vivid, concrete visual details. Add setting, lighting, textures, or objects. " f"Do NOT repeat. " f"{avoid_clause}" f"Never use storytelling, feelings, or narrative context." f"Your output must be suitable for direct input into an AI image generation model." f"Give a new idea from the previous" f"{last_output}" ) else: user_message_content = ( f"DO NOT explain your reasoning or describe what you are doing." f"Improve the following prompt for visual clarity and composition. " f"Structure for better flow, but stay under {max_tokens} tokens. " f"{avoid_clause}" f"Never use storytelling, feelings, or narrative context." f"Your output must be suitable for direct input into an AI image generation model." f"Give a new idea from the previous" f"{last_output}" ) system_message_content += avoid_clause if randomize_temp: llm_temp = random.uniform(0.1, 1.0) raw_result = "" if is_ollama_model: api_url = "http://localhost:11434/api/generate" if randomize_seed: llm_seed = random.randint(0, 1000000000) payload = { "model": actual_model_name, "prompt": system_message_content + "\n" + user_message_content, "stream": False, "options": { "temperature": llm_temp, "seed": llm_seed, } } response = requests.post( api_url, json=payload, timeout=PREDICTION_TIMEOUT ) response.raise_for_status() raw_result = response.json().get("response", "").strip() elif is_lmstudio_model: if lm_studio_llm_instance is None: raise ConnectionError("LM Studio LLM instance not initialized for API call.") chat = lms.Chat() chat.add_user_message(system_message_content + "\n" + user_message_content) stream = lm_studio_llm_instance.respond_stream(chat, config={"repeatPenalty":1.1, "temperature":llm_temp, "seed": llm_seed,}, on_message=chat.append) cancelled = False def timeout_handler(): nonlocal cancelled cancelled = True stream.cancel() timer = threading.Timer(PREDICTION_TIMEOUT, timeout_handler) timer.start() try: for _ in stream: if cancelled: break if not cancelled: result_obj = stream.result() if hasattr(result_obj, "text"): raw_result = result_obj.text.strip() else: raw_result = str(result_obj).strip() else: raise TimeoutError("LM Studio response generation cancelled due to timeout.") finally: timer.cancel() token_count = estimate_tokens(raw_result) print(f"Idea: {idx + 1} Attempt: {attempt}/{max_attempts} LLM result: {raw_result}") print(f"→ Token count: {token_count} (target: {min_tokens}–{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 timeout_number = 0 time.sleep(0.5) continue if min_tokens <= token_count <= max_tokens: used_phrases.append(raw_result) print("✔️ Prompt accepted.") timeout_number = 0 if keepllm == False: if is_ollama_model: clear_ollama_model() elif is_lmstudio_model and lm_studio_llm_instance: model = lms.llm() model.unload() generated_prompts.append(raw_result) break last_output = raw_result print(f"⚠️ Prompt out of bounds. Retrying...\n") timeout_number = 0 time.sleep(0.5) except requests.exceptions.HTTPError as e: if e.response.status_code == 500: print(f"⚠️ Attempt {attempt}/{max_attempts} HTTP 500 error. Retrying...\n") if keepllm == False: if is_ollama_model: clear_ollama_model() time.sleep(1) continue else: error_msg = f"[LLM API Error] HTTP Error {e.response.status_code}: {e.response.text}" print(error_msg) return (error_msg, negative, idea) except requests.exceptions.Timeout: if timeout_number < 5: print(f"⚠️ Attempt {attempt}/{max_attempts} timed out. Retrying...\n") time.sleep(1) timeout_number += 1 continue else: error_msg = f"Too many timeouts (5). Falling back to idea." print(f"❌ {error_msg}") generated_prompts.append(sub_idea) timeout_number = 0 if keepllm == False: if is_ollama_model: clear_ollama_model() elif is_lmstudio_model and lm_studio_llm_instance: model = lms.llm() model.unload() break except TimeoutError as e: if timeout_number < 5: print(f"⚠️ Attempt {attempt}/{max_attempts} LM Studio stream timeout. Retrying...\n") time.sleep(1) timeout_number += 1 continue else: error_msg = f"Too many LM Studio timeouts (5). Falling back to idea." print(f"❌ {error_msg}") generated_prompts.append(sub_idea) timeout_number = 0 if keepllm == False: if is_lmstudio_model and lm_studio_llm_instance: model = lms.llm() model.unload() break except Exception as e: error_msg = f"[LLM Error] {str(e)}" print(error_msg) if keepllm == False: if is_ollama_model: clear_ollama_model() elif is_lmstudio_model and lm_studio_llm_instance: model = lms.llm() model.unload() return (error_msg, negative, idea) else: print(f"❌ Max attempts for idea '{sub_idea}' reached. Using original as fallback.") timeout_number = 0 if keepllm == False: if is_ollama_model: clear_ollama_model() elif is_lmstudio_model and lm_studio_llm_instance: model = lms.llm() model.unload() generated_prompts.append(sub_idea) final_prompt = " BREAK ".join(generated_prompts) print(f"\nFinal Generated Prompt: {final_prompt}") pbar.update_absolute(max_attempts) try: with open(prompt_log_file, "w+", encoding="utf-8") as f: f.write(final_prompt) except Exception as log_error: print(f"⚠️ Failed to write prompt to file: {log_error}") return (final_prompt, negative, idea) def ui(self, inputs, outputs): prompt_str = outputs[0] if isinstance(outputs, (list, tuple)) and outputs else "" return { "prompt": f"🧠 Generated Prompt:\n{prompt_str}" } NODE_CLASS_MAPPINGS = { "OllamaPromptFromIdea": OllamaPromptFromIdea, } NODE_DISPLAY_NAME_MAPPINGS = { "OllamaPromptFromIdea": "🧠 Ollama & LM Studio Prompt From Idea", }