- Add DeepSeek Prompt Rewriter node for fast API-based rewriting - Add Load Image with Prompt node for metadata extraction - Add config.json.example with API backend templates - Add flux2/schema.yaml - Update .gitignore to exclude dev/ and config.json (contains API keys)
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
|
|
import os
|
|
import json
|
|
from PIL import Image, ImageOps
|
|
import torch
|
|
import numpy as np
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
folder_paths = None
|
|
|
|
class EricLoadImageWithPrompt:
|
|
"""
|
|
Loads an image and extracts generation metadata (prompt) if available.
|
|
Supports standard ComfyUI image loading + metadata extraction.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
|
return {
|
|
"required": {
|
|
"image": (sorted(files), {"image_upload": True})
|
|
}
|
|
}
|
|
|
|
CATEGORY = "Eric Prompt Enhancers"
|
|
RETURN_TYPES = ("IMAGE", "MASK", "STRING")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "PROMPT_TEXT")
|
|
FUNCTION = "load_image"
|
|
|
|
def load_image(self, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
|
|
# Process image for ComfyUI (Tensor)
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
# Create mask
|
|
if 'A' in i.getbands():
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
mask = 1. - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
|
|
|
# Extract Metadata
|
|
prompt_text = ""
|
|
|
|
# 1. Try standard PNG info (parameters/Comment)
|
|
# A1111 stores generation info in 'parameters'
|
|
if i.info:
|
|
if 'parameters' in i.info:
|
|
prompt_text = i.info['parameters']
|
|
elif 'Comment' in i.info:
|
|
# Some software uses Comment
|
|
try:
|
|
# Try to parse JSON if it looks like it
|
|
comment = i.info['Comment']
|
|
if comment.strip().startswith('{'):
|
|
data = json.loads(comment)
|
|
prompt_text = data.get('prompt', comment)
|
|
else:
|
|
prompt_text = comment
|
|
except:
|
|
prompt_text = i.info['Comment']
|
|
|
|
# 2. Try ComfyUI workflow metadata (often in 'workflow' or 'prompt')
|
|
# This is complex because it's the whole graph. We try to find a primitive string or positive prompt.
|
|
# For now, we'll stick to the text-based 'parameters' which is the standard for sharing prompts.
|
|
|
|
# Clean up A1111 'parameters' block to get just the positive prompt
|
|
# Usually it's: "Positive Prompt\nNegative Prompt: ...\nSteps: ..."
|
|
if prompt_text and "Negative prompt:" in prompt_text:
|
|
prompt_text = prompt_text.split("Negative prompt:")[0].strip()
|
|
elif prompt_text and "Steps:" in prompt_text:
|
|
prompt_text = prompt_text.split("Steps:")[0].strip()
|
|
|
|
return (image, mask, prompt_text)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, image):
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
m = os.path.getmtime(image_path)
|
|
return m
|