356 lines
18 KiB
Python
356 lines
18 KiB
Python
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/local_prompt_generator.py
|
||
import hashlib
|
||
import os
|
||
import json
|
||
|
||
# --- Simple logging utility for this package ---
|
||
def log(message):
|
||
print(f"[LocalPromptGen] {message}")
|
||
|
||
# --- Replicate necessary preset logic locally ---
|
||
# We need to replicate the core preset loading/management logic here.
|
||
|
||
# --- Define paths and functions for user presets ---
|
||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||
USER_PRESETS_FILE = os.path.join(script_directory, "user_kontext_presets.json")
|
||
|
||
# --- Replicate Built-in Presets (from MieNodes' prompt_generator.py) ---
|
||
KONTEXT_PRESETS = {
|
||
"Kid Clothes": {
|
||
"system": (
|
||
"You are an expert prompt generator for Flux Kontext. Your task is to place clothing described by the user onto a baby or child with Arab facial features, while preserving all material, color, and style details from the original clothing description. "
|
||
"Input: "
|
||
"1. [IMAGE DESC]: A clean and specific description of a piece of baby or children’s clothing. "
|
||
"2. [USER INTENT]: Typically vague, like 'put on baby' or 'make it worn'. "
|
||
"Rules: "
|
||
"Output only a single Flux prompt (no extra labels, quotes, or explanation). "
|
||
"Always apply the clothes on a realistic Arab baby or toddler (chubby cheeks, warm skin tone, soft features, natural lighting). "
|
||
"Retain all textures, cuts, stitching, patterns, and colors of the clothing exactly. "
|
||
"Set the baby in a realistic environment: white blanket, nursery, soft studio setting. "
|
||
"Describe lighting: soft diffused lighting, warm tones, shallow depth of field. "
|
||
"Use fashion product photography terms and camera specs. "
|
||
"Example: "
|
||
"[IMAGE DESC]: A pale blue cotton baby romper with white cloud patterns, soft buttons and footies. "
|
||
"[USER INTENT]: put it on a baby "
|
||
"Output: "
|
||
"A photorealistic Arab baby with round cheeks and light tan skin, lying on a soft white fleece blanket, wearing a pale blue cotton romper with stitched white cloud patterns and integrated footies, captured in soft natural lighting with shallow DOF, cozy nursery ambiance, 8K resolution, cinematic baby fashion catalog style."
|
||
)
|
||
},
|
||
"Women’s Clothes": {
|
||
"system": (
|
||
"You are an expert visual prompt composer for Flux Kontext. Your job is to place the clothing described by the user onto a realistic Muslim woman model in her 20s with Arab features, while preserving all original garment details. "
|
||
"Input: "
|
||
"1. [IMAGE DESC]: A detailed visual description of a clothing item (e.g., abaya, blouse, scarf). "
|
||
"2. [USER INTENT]: Often short, like 'make her wear it' or 'put on a model'. "
|
||
"Rules: "
|
||
"Output a single photorealistic prompt. "
|
||
"Always depict a modest Arab Muslim woman, 20s, wearing the clothes naturally and confidently. "
|
||
"Keep all design details: fabric texture, color, embroidery, shape. "
|
||
"Style should reflect high-end modest fashion or Islamic wear catalog. "
|
||
"Scene: outdoor urban, soft studio light, clean modern backdrop, or lifestyle setting. "
|
||
"Describe lighting, pose, camera angle, clothing drape and body fit. "
|
||
"Example: "
|
||
"[IMAGE DESC]: A beige linen abaya with gold-stitched cuffs and a waist tie. "
|
||
"[USER INTENT]: put it on muslim woman "
|
||
"Output: "
|
||
"A young Arab Muslim woman in her 20s with soft olive skin and delicate features, wearing a beige linen abaya with subtle gold-stitched cuffs and a loosely tied waist sash, standing in a minimalist white studio under soft diffused lighting, side-facing pose with natural folds in the fabric, fashion editorial style with cinematic shallow DOF."
|
||
)
|
||
},
|
||
"Beauty Product Use": {
|
||
"system": (
|
||
"You are a creative image prompt generator for Flux Kontext focused on health, skincare, and beauty products. Your job is to place the described product in an artistic or commercial setting — either used by a model (Arab man or woman), or placed within a creative, beauty-inspired scene — while preserving all visual and material details of the product. "
|
||
"Input: "
|
||
"1. [IMAGE DESC]: A detailed image or design of a product (e.g., oil bottle, cream jar, serum tube). "
|
||
"2. [USER INTENT]: Usually vague like 'make woman use it' or 'put it in beautiful place'. "
|
||
"Rules: "
|
||
"Output a single detailed visual prompt. "
|
||
"Product must appear clearly: keep all branding, color, bottle design intact. "
|
||
"If used: describe the Arab model using it naturally (e.g., applying to face, holding dropper). "
|
||
"If placed: build a beautiful, creative background — spa setting, nature elements, mirror, floral. "
|
||
"Use keywords: soft light, high-end spa, cinematic DOF, luxury textures, minimalist props. "
|
||
"Example: "
|
||
"[IMAGE DESC]: A small amber glass dropper bottle with golden label and black cap. "
|
||
"[USER INTENT]: woman use it "
|
||
"Output: "
|
||
"An elegant Arab woman with natural tan skin and voluminous dark hair, holding an amber glass dropper bottle with a gold label and black cap, delicately applying serum to her cheek in front of a glowing vanity mirror, soft candle-lit ambiance, subtle shadows, cinematic close-up, bokeh background, luxury skincare ad style."
|
||
)
|
||
},
|
||
"Beauty Product Display": {
|
||
"system": (
|
||
"You are an expert in luxury beauty product visuals for Flux Kontext. Place the described product in a stunning, creative, brand-inspired environment (without any models), using beauty aesthetics and visual storytelling. "
|
||
"Input: "
|
||
"1. [IMAGE DESC]: Description of a product container or design (e.g., shampoo bottle, bar soap, lip balm). "
|
||
"2. [USER INTENT]: Usually vague like 'put it in nature' or 'make it pretty'. "
|
||
"Rules: "
|
||
"Output only one final prompt. "
|
||
"Focus on creative display — natural materials, elegant composition, light play. "
|
||
"Use props like glass trays, greenery, water droplets, stone, sand, soft cloth. "
|
||
"Mention reflections, shadows, lighting temperature. "
|
||
"Product must stand out, sharply rendered. "
|
||
"Example: "
|
||
"[IMAGE DESC]: A matte pink soap bar with embossed logo. "
|
||
"[USER INTENT]: put in creative background "
|
||
"Output: "
|
||
"A matte pink soap bar with embossed luxury logo resting on a smooth white marble tray, surrounded by pale rose petals and water droplets, gentle morning light filtering through a frosted glass window, soft shadows, minimalist spa aesthetic, 8K beauty product photography."
|
||
)
|
||
},
|
||
}
|
||
# --- End of KONTEXT_PRESETS ---
|
||
|
||
def load_user_presets():
|
||
"""Load user-defined presets from the JSON file."""
|
||
if os.path.exists(USER_PRESETS_FILE):
|
||
try:
|
||
with open(USER_PRESETS_FILE, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except (json.JSONDecodeError, IOError) as e:
|
||
log(f"Error loading user presets: {e}")
|
||
return {}
|
||
return {}
|
||
|
||
def get_all_kontext_presets():
|
||
"""Get a combined dictionary of built-in and user presets."""
|
||
# Start with built-in presets
|
||
all_presets = KONTEXT_PRESETS.copy()
|
||
# Load and add user presets, potentially overriding built-ins if names clash (usually desired)
|
||
user_presets = load_user_presets()
|
||
all_presets.update(user_presets)
|
||
return all_presets
|
||
|
||
# Define category directly for this package
|
||
LOCAL_PROMPT_CATEGORY = "Local LLM Nodes/Prompt Generators"
|
||
|
||
class LocalKontextPromptGenerator(object):
|
||
"""
|
||
A version of KontextPromptGenerator designed to work specifically with
|
||
the LocalLLMServiceConnector. Replicates the core prompt generation logic exactly.
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
# --- Replicate INPUT_TYPES exactly ---
|
||
all_presets = get_all_kontext_presets()
|
||
# Ensure there's a default if the list is somehow empty
|
||
default_preset_key = next(iter(all_presets.keys()), "") if all_presets.keys() else ""
|
||
return {
|
||
"required": {
|
||
# --- CRITICAL: Use the standard LLM_SERVICE_CONNECTOR type ---
|
||
# This matches the RETURN_TYPES of SetLocalLLMServiceConnector
|
||
# and is the same type expected by the original KontextPromptGenerator.
|
||
"llm_service_connector": ("LLMServiceConnector",),
|
||
"image1_description": ("STRING", {"default": "", "multiline": True, "tooltip": "Describe the first image"}),
|
||
"image2_description": ("STRING", {"default": "", "multiline": True, "tooltip": "Describe the second image"}),
|
||
"edit_instruction": ("STRING", {"default": "", "multiline": True}),
|
||
"preset": (list(all_presets.keys()), {"default": default_preset_key}),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True,
|
||
"tooltip": "The random seed used for creating the noise."}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("kontext_prompt",)
|
||
FUNCTION = "generate_kontext_prompt"
|
||
CATEGORY = LOCAL_PROMPT_CATEGORY
|
||
|
||
# --- Replicate the core generate_kontext_prompt method exactly ---
|
||
def generate_kontext_prompt(self, llm_service_connector, image1_description, image2_description, edit_instruction, preset, seed=None):
|
||
"""
|
||
Replicates the exact core logic from KontextPromptGenerator.generate_kontext_prompt.
|
||
"""
|
||
# --- Exact replication of the original logic ---
|
||
all_presets = get_all_kontext_presets()
|
||
preset_data = all_presets.get(preset)
|
||
# --- Fixed Syntax Error ---
|
||
if not preset_data: # <-- Corrected condition check
|
||
raise ValueError(f"Unknown preset: {preset}")
|
||
|
||
# --- Refined user content construction for Flux prompt generation ---
|
||
# --- Refined user content construction for strict interpretation ---
|
||
def safe_str_convert(value):
|
||
"""Converts input to string, handling lists and None."""
|
||
if value is None:
|
||
return ""
|
||
if isinstance(value, list):
|
||
return " ".join(str(item) for item in value if item is not None)
|
||
return str(value)
|
||
|
||
# Get and convert inputs
|
||
raw_img1_desc = image1_description
|
||
# raw_img2_desc = image2_description # Explicitly ignored for now
|
||
raw_edit_inst = edit_instruction
|
||
|
||
img1_desc_str = safe_str_convert(raw_img1_desc).strip()
|
||
# img2_desc_str = safe_str_convert(raw_img2_desc).strip() # Ignored
|
||
edit_inst_str = safe_str_convert(raw_edit_inst).strip()
|
||
|
||
# --- Key Change: Sharper distinction and labeling ---
|
||
# Construct user_content with very clear, labeled sections
|
||
user_content_parts = []
|
||
|
||
if img1_desc_str:
|
||
# Present the base description as the subject/context
|
||
user_content_parts.append(f"IMAGE 1 DESCRIPTION: {img1_desc_str}")
|
||
# user_content_parts.append(f"[SUBJECT]: {img1_desc_str}")
|
||
|
||
if edit_inst_str:
|
||
# Directly frame edits as Flux adjustments
|
||
user_content_parts.append(f"FLUX EDIT INSTRUCTION: {edit_inst_str}")
|
||
else:
|
||
# Default to a Flux-style improvement request
|
||
user_content_parts.append("FLUX EDIT INSTRUCTION: Enhance details and realism.")
|
||
|
||
# Combine parts
|
||
user_content = " ".join(user_content_parts) # Simple space separator
|
||
|
||
# Fallback if somehow inputs were empty
|
||
if not user_content.strip():
|
||
user_content = "IMAGE 1 DESCRIPTION: A product on a white background. FLUX EDIT INSTRUCTION: Make it photorealistic."
|
||
|
||
# --- End of Refined user content construction ---
|
||
|
||
# Rest of the method remains the same...
|
||
messages = [
|
||
{"role": "system", "content": preset_data["system"]},
|
||
{"role": "user", "content": user_content},
|
||
]
|
||
|
||
# --- Key Change: Use the local connector's invoke method ---
|
||
# Pass the messages and seed.
|
||
try:
|
||
kontext_prompt = llm_service_connector.invoke(messages, seed=seed)
|
||
# Ensure the output is stripped, like the original
|
||
return (kontext_prompt.strip(),)
|
||
except Exception as e:
|
||
# Handle potential errors during local LLM invocation
|
||
error_msg = f"Error generating kontext prompt with local LLM: {str(e)}"
|
||
log(error_msg) # Log to console (using the 'log' function defined in this file)
|
||
# Return the error message as the prompt string so the workflow doesn't crash silently
|
||
return (error_msg,)
|
||
|
||
|
||
# --- Replicate the is_changed method for proper caching ---
|
||
def is_changed(self, llm_service_connector, image1_description, image2_description, edit_instruction, preset, seed):
|
||
"""
|
||
Replicates the exact core logic from KontextPromptGenerator.is_changed.
|
||
"""
|
||
try:
|
||
hasher = hashlib.md5()
|
||
hasher.update(image1_description.encode('utf-8'))
|
||
hasher.update(image2_description.encode('utf-8'))
|
||
hasher.update(edit_instruction.encode('utf-8'))
|
||
hasher.update(preset.encode('utf-8'))
|
||
hasher.update(str(seed).encode('utf-8'))
|
||
|
||
# Incorporate preset system prompt
|
||
all_presets = get_all_kontext_presets()
|
||
preset_data = all_presets.get(preset)
|
||
# --- Fixed Syntax Error ---
|
||
if preset_data and "system" in preset_data: # <-- Corrected condition and variable name
|
||
hasher.update(preset_data["system"].encode('utf-8'))
|
||
|
||
# Incorporate connector state (replicate original logic)
|
||
# Original tries get_state(), then falls back to specific attributes.
|
||
# For a local connector, str(connector) is suitable.
|
||
try:
|
||
# Try get_state if it exists on the local connector (unlikely for our simple one)
|
||
connector_state = llm_service_connector.get_state()
|
||
except AttributeError:
|
||
# Fallback: Use a string representation of the connector object
|
||
# This works if the connector's identity is tied to its instance/model path.
|
||
connector_state = str(llm_service_connector)
|
||
|
||
hasher.update(connector_state.encode('utf-8'))
|
||
|
||
return hasher.hexdigest()
|
||
except Exception as e:
|
||
# If hashing fails, force re-run
|
||
log(f"is_changed error: {e}")
|
||
return float("nan") # Always re-run
|
||
|
||
# --- Replicate the AddUserKontextPreset and RemoveUserKontextPreset classes ---
|
||
# These manage the local user presets file.
|
||
|
||
class AddUserLocalKontextPreset: # <-- Changed class name prefix
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"preset_name": ("STRING", {"default": ""}),
|
||
"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("BOOLEAN", "STRING")
|
||
RETURN_NAMES = ("success", "log")
|
||
FUNCTION = "add_preset"
|
||
CATEGORY = LOCAL_PROMPT_CATEGORY
|
||
|
||
def add_preset(self, preset_name, system_prompt):
|
||
import datetime
|
||
if not preset_name or not system_prompt:
|
||
log = "Preset name and system prompt must not be empty."
|
||
return (False, log)
|
||
|
||
user_presets = load_user_presets()
|
||
if preset_name in user_presets:
|
||
log = f"Preset '{preset_name}' already exists (custom preset)."
|
||
return (False, log)
|
||
|
||
user_presets[preset_name] = {"system": system_prompt}
|
||
# Save to the local user presets file
|
||
try:
|
||
with open(USER_PRESETS_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(user_presets, f, ensure_ascii=False, indent=2)
|
||
now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
log = f"Preset '{preset_name}' added successfully at {now}."
|
||
return (True, log)
|
||
except Exception as e:
|
||
log = f"Error saving preset: {e}"
|
||
return (False, log)
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, **kwargs):
|
||
# Always re-run when called to check for file changes
|
||
return float("nan")
|
||
|
||
class RemoveUserLocalKontextPreset: # <-- Changed class name prefix
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
# Load current presets to populate the dropdown
|
||
user_presets = load_user_presets()
|
||
preset_names = list(user_presets.keys())
|
||
# Provide a default if the list is empty
|
||
default_name = preset_names[0] if preset_names else ""
|
||
return {
|
||
"required": {
|
||
"preset_name": (preset_names, {"default": default_name}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("BOOLEAN", "STRING")
|
||
RETURN_NAMES = ("success", "log")
|
||
FUNCTION = "remove_preset"
|
||
CATEGORY = LOCAL_PROMPT_CATEGORY
|
||
|
||
def remove_preset(self, preset_name):
|
||
import datetime
|
||
user_presets = load_user_presets()
|
||
if preset_name in user_presets:
|
||
del user_presets[preset_name]
|
||
try:
|
||
# Save to the local user presets file
|
||
with open(USER_PRESETS_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(user_presets, f, ensure_ascii=False, indent=2)
|
||
now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
log = f"Preset '{preset_name}' removed successfully at {now}."
|
||
return (True, log)
|
||
except Exception as e:
|
||
log = f"Error saving after removal: {e}"
|
||
return (False, log)
|
||
else:
|
||
log = f"Preset '{preset_name}' not found in user presets."
|
||
return (False, log)
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, **kwargs):
|
||
# Always re-run when called to check for file changes
|
||
return float("nan") |