diff --git a/nodes.py b/nodes.py index c5ac3bb..9c18ae3 100644 --- a/nodes.py +++ b/nodes.py @@ -1,14 +1,12 @@ - import os import json import urllib.request import urllib.error from dotenv import load_dotenv -# Load variables from .env if it exists load_dotenv() -class BritaPromptGenerator: +class PromptGeneratorCore: def __init__(self): self.api_key = os.getenv("OPENAI_API_KEY") if not self.api_key: @@ -17,26 +15,8 @@ class BritaPromptGenerator: self.model = "chatgpt-4o-latest" self.temperature = 0.88 - self.system_prompt = ( - "You are an expert prompt engineer for AI image generation models, specifically for the 'Flux' model. " - "Your task is to generate creative and diverse prompts for 'Flux' model. Each prompt must include the Lora trigger: " - "'brita water filter with blue lid, product photography'. The scene should feature a cute animal " - "in a living room, curiously observing this water filter. Ensure variety in animal type, " - "living room style, lighting, and the animal's specific curious action. " - "Output exactly Flux prompts directly and based on best practice. Format the output as line separated plain text without irrelevant details." - ) - - self.user_prompt = ( - "Generate 25 image generation prompts based on the following core elements:\n" - "1. Model: Flux (implying detailed, high-quality output desired)\n" - "2. Subject: A cute animal (e.g., kitten, puppy, bunny, hamster, small fox, baby owl).\n" - "3. Object: A Brita water filter pitcher.\n" - "4. Lora Trigger: Must include 'brita water filter with blue lid, product photography'.\n" - "5. Setting: A living room (e.g., modern, cozy, minimalist, bohemian).\n" - "6. Interaction: The animal is curious about the water filter (e.g., sniffing, pawing, tilting head, peering into it).\n" - "7. Style: Aim for photorealistic or beautifully illustrative, with good lighting.\n" - "Provide the output should start with the LoRA trigger. Response should be in line separated plain text without any irrelevant details." - ) + self.system_prompt = "You are a helpful assistant that creates prompts for AI applications." + self.user_prompt = "Generate 10 high-quality prompts for testing purposes." def generate(self): headers = { @@ -74,7 +54,7 @@ class Node: CATEGORY = "flux/prompt" def __init__(self): - self.generator = BritaPromptGenerator() + self.generator = PromptGeneratorCore() @classmethod def INPUT_TYPES(cls): @@ -85,7 +65,7 @@ class Node: FUNCTION = "generate_prompts" OUTPUT_NODE = False - DESCRIPTION = "Generate 25 Flux prompts with Brita water filter and cute animal interaction." + DESCRIPTION = "Generate AI prompts using OpenAI Chat API." def generate_prompts(self): output = self.generator.generate() @@ -93,9 +73,9 @@ class Node: NODE_CLASS_MAPPINGS = { - "BritaPromptGeneratorNode": Node + "PromptGeneratorNode": Node } NODE_DISPLAY_NAME_MAPPINGS = { - "BritaPromptGeneratorNode": "Brita Prompt Generator" + "PromptGeneratorNode": "Prompt Generator" }