Enhance SaveImage functionality with custom prompt support

- Added a new optional parameter `custom_prompt` to the SaveImage class methods to allow users to override the default prompt.
- Updated the `format_metadata` method to utilize the custom prompt if provided.
- Modified the `save_images` and `process_image` methods to accept and pass the custom prompt through the workflow processing.
This commit is contained in:
Will Miao
2025-04-04 07:47:46 +08:00
parent 6af12d1acc
commit 9613199152
2 changed files with 15 additions and 6 deletions

View File

@@ -34,6 +34,7 @@ class SaveImage:
"file_format": (["png", "jpeg", "webp"],), "file_format": (["png", "jpeg", "webp"],),
}, },
"optional": { "optional": {
"custom_prompt": ("STRING", {"default": "", "forceInput": True}),
"lossless_webp": ("BOOLEAN", {"default": True}), "lossless_webp": ("BOOLEAN", {"default": True}),
"quality": ("INT", {"default": 100, "min": 1, "max": 100}), "quality": ("INT", {"default": 100, "min": 1, "max": 100}),
"embed_workflow": ("BOOLEAN", {"default": False}), "embed_workflow": ("BOOLEAN", {"default": False}),
@@ -60,7 +61,7 @@ class SaveImage:
return item.get('sha256') return item.get('sha256')
return None return None
async def format_metadata(self, parsed_workflow): async def format_metadata(self, parsed_workflow, custom_prompt=None):
"""Format metadata in the requested format similar to userComment example""" """Format metadata in the requested format similar to userComment example"""
if not parsed_workflow: if not parsed_workflow:
return "" return ""
@@ -69,6 +70,10 @@ class SaveImage:
prompt = parsed_workflow.get('prompt', '') prompt = parsed_workflow.get('prompt', '')
negative_prompt = parsed_workflow.get('negative_prompt', '') negative_prompt = parsed_workflow.get('negative_prompt', '')
# Override prompt with custom_prompt if provided
if custom_prompt:
prompt = custom_prompt
# Extract loras from the prompt if present # Extract loras from the prompt if present
loras_text = parsed_workflow.get('loras', '') loras_text = parsed_workflow.get('loras', '')
lora_hashes = {} lora_hashes = {}
@@ -240,7 +245,8 @@ class SaveImage:
return filename return filename
def save_images(self, images, filename_prefix, file_format, prompt=None, extra_pnginfo=None, def save_images(self, images, filename_prefix, file_format, prompt=None, extra_pnginfo=None,
lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True): lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True,
custom_prompt=None):
"""Save images with metadata""" """Save images with metadata"""
results = [] results = []
@@ -248,11 +254,12 @@ class SaveImage:
parser = WorkflowParser() parser = WorkflowParser()
if prompt: if prompt:
parsed_workflow = parser.parse_workflow(prompt) parsed_workflow = parser.parse_workflow(prompt)
print("parsed_workflow", parsed_workflow)
else: else:
parsed_workflow = {} parsed_workflow = {}
# Get or create metadata asynchronously # Get or create metadata asynchronously
metadata = asyncio.run(self.format_metadata(parsed_workflow)) metadata = asyncio.run(self.format_metadata(parsed_workflow, custom_prompt))
# Process filename_prefix with pattern substitution # Process filename_prefix with pattern substitution
filename_prefix = self.format_filename(filename_prefix, parsed_workflow) filename_prefix = self.format_filename(filename_prefix, parsed_workflow)
@@ -338,7 +345,8 @@ class SaveImage:
return results return results
def process_image(self, images, filename_prefix="ComfyUI", file_format="png", prompt=None, extra_pnginfo=None, def process_image(self, images, filename_prefix="ComfyUI", file_format="png", prompt=None, extra_pnginfo=None,
lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True): lossless_webp=True, quality=100, embed_workflow=False, add_counter_to_filename=True,
custom_prompt=""):
"""Process and save image with metadata""" """Process and save image with metadata"""
# Make sure the output directory exists # Make sure the output directory exists
os.makedirs(self.output_dir, exist_ok=True) os.makedirs(self.output_dir, exist_ok=True)
@@ -356,7 +364,8 @@ class SaveImage:
lossless_webp, lossless_webp,
quality, quality,
embed_workflow, embed_workflow,
add_counter_to_filename add_counter_to_filename,
custom_prompt if custom_prompt.strip() else None
) )
return (images,) return (images,)

View File

@@ -134,7 +134,7 @@ def transform_lora_loader(inputs: Dict) -> Dict:
"loras": " ".join(lora_texts) "loras": " ".join(lora_texts)
} }
if "clip" in inputs: if "clip" in inputs and isinstance(inputs["clip"], dict):
result["clip_skip"] = inputs["clip"].get("clip_skip", "-1") result["clip_skip"] = inputs["clip"].get("clip_skip", "-1")
return result return result