AEmotionStudioandClaude Opus 4.5
36c5d06cf8
The refactored prompt extractor only detected CLIPTextEncode nodes, losing support for SDXL-specific nodes that was present in the original code. This caused SDXL workflows using SDXLPromptEncoder or SDXLTextEncode nodes to not have their prompts included in Discord messages. Changes: - Add PROMPT_NODE_TYPES list with CLIPTextEncode, SDXLPromptEncoder, and SDXLTextEncode - Rename functions from clip-specific to generic prompt node naming - Add fallback detection for custom text/encode/prompt nodes - Update docstrings to reflect broader node type support Fixes regression identified in PR #35 review. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
271 lines
9.5 KiB
Python
271 lines
9.5 KiB
Python
"""
|
|
Prompt Extraction Utilities for ComfyUI-DiscordSend
|
|
|
|
Extracts positive and negative prompts from ComfyUI workflow data.
|
|
"""
|
|
|
|
import json
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
|
|
# Common negative prompt indicators
|
|
NEGATIVE_INDICATORS = [
|
|
"bad quality", "deformed", "blurry", "low quality", "worst quality",
|
|
"ugly", "disfigured", "low res", "poorly drawn", "mutation",
|
|
"extra limbs", "bad anatomy", "watermark", "text", "signature"
|
|
]
|
|
|
|
# Node types that can contain prompts
|
|
PROMPT_NODE_TYPES = [
|
|
"CLIPTextEncode", # Standard SD 1.5 prompt node
|
|
"SDXLPromptEncoder", # SDXL prompt encoder
|
|
"SDXLTextEncode", # Another SDXL text node
|
|
]
|
|
|
|
|
|
def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Optional[str]]:
|
|
"""
|
|
Extract positive and negative prompts from workflow data.
|
|
|
|
Analyzes ComfyUI workflow structure to find prompt nodes (CLIPTextEncode,
|
|
SDXLPromptEncoder, SDXLTextEncode, and other text encoding nodes) and
|
|
determine which contains the positive vs negative prompt.
|
|
|
|
Args:
|
|
workflow_data: The workflow data dictionary or JSON string
|
|
|
|
Returns:
|
|
A tuple of (positive_prompt, negative_prompt) or (None, None) if not found
|
|
"""
|
|
if workflow_data is None:
|
|
return None, None
|
|
|
|
# Convert string to dict if necessary
|
|
if isinstance(workflow_data, str):
|
|
try:
|
|
data = json.loads(workflow_data)
|
|
except json.JSONDecodeError:
|
|
return None, None
|
|
else:
|
|
data = workflow_data
|
|
|
|
if not isinstance(data, dict):
|
|
return None, None
|
|
|
|
positive_prompt = None
|
|
negative_prompt = None
|
|
|
|
# Find prompt nodes (CLIPTextEncode, SDXL nodes, etc.)
|
|
if "nodes" in data:
|
|
nodes = data["nodes"]
|
|
else:
|
|
# Check if it's API format (dict of nodes)
|
|
nodes = data
|
|
|
|
prompt_nodes = _find_prompt_nodes(nodes)
|
|
|
|
if not prompt_nodes:
|
|
return None, None
|
|
|
|
# Determine positive/negative based on content and structure
|
|
if len(prompt_nodes) == 1:
|
|
# Single prompt node - assume it's the positive prompt
|
|
positive_prompt = _get_prompt_text(prompt_nodes[0])
|
|
elif len(prompt_nodes) >= 2:
|
|
# Multiple prompt nodes - need to determine which is which
|
|
positive_prompt, negative_prompt = _classify_prompts(prompt_nodes, data)
|
|
|
|
# Ensure we return empty string for negative if we have positive but not negative
|
|
if positive_prompt is not None and negative_prompt is None:
|
|
negative_prompt = ""
|
|
|
|
return positive_prompt, negative_prompt
|
|
|
|
|
|
def _find_prompt_nodes(nodes: Union[List, Dict]) -> List[Dict]:
|
|
"""Find all prompt nodes (CLIPTextEncode, SDXL nodes, etc.) in the workflow."""
|
|
prompt_nodes = []
|
|
|
|
if isinstance(nodes, list):
|
|
for node in nodes:
|
|
if _is_prompt_node(node):
|
|
prompt_nodes.append(node)
|
|
elif isinstance(nodes, dict):
|
|
for node_id, node in nodes.items():
|
|
if _is_prompt_node(node):
|
|
node_copy = dict(node)
|
|
node_copy["id"] = node_id
|
|
prompt_nodes.append(node_copy)
|
|
|
|
return prompt_nodes
|
|
|
|
|
|
def _is_prompt_node(node: Any) -> bool:
|
|
"""Check if a node is a prompt node (CLIPTextEncode, SDXL, etc.) with valid text."""
|
|
if not isinstance(node, dict):
|
|
return False
|
|
|
|
# Handle both Workflow format (type) and API format (class_type)
|
|
node_type = node.get("type") or node.get("class_type")
|
|
|
|
# Check against known prompt node types
|
|
if node_type not in PROMPT_NODE_TYPES:
|
|
# Also check for dynamic text/prompt nodes (e.g., custom nodes)
|
|
if node_type and ("Text" in node_type and ("Encode" in node_type or "Prompt" in node_type)):
|
|
pass # Allow these through
|
|
else:
|
|
return False
|
|
|
|
# Check for text in either widgets_values (Workflow) or inputs (API)
|
|
text = _get_prompt_text(node)
|
|
return text is not None
|
|
|
|
|
|
def _get_prompt_text(node: Dict) -> Optional[str]:
|
|
"""Extract the prompt text from a prompt node."""
|
|
# Workflow format (widgets_values)
|
|
widgets = node.get("widgets_values", [])
|
|
if isinstance(widgets, list) and len(widgets) > 0 and isinstance(widgets[0], str):
|
|
return widgets[0]
|
|
|
|
# API format (inputs.text)
|
|
inputs = node.get("inputs", {})
|
|
if isinstance(inputs, dict) and isinstance(inputs.get("text"), str):
|
|
return inputs.get("text")
|
|
|
|
return None
|
|
|
|
|
|
def _classify_prompts(prompt_nodes: List[Dict], workflow_data: Dict) -> Tuple[Optional[str], Optional[str]]:
|
|
"""
|
|
Classify which prompt nodes contain positive vs negative prompts.
|
|
|
|
Uses multiple heuristics:
|
|
1. Content analysis (negative prompts often contain quality-related terms)
|
|
2. Connection analysis (traces connections to sampler nodes)
|
|
"""
|
|
if not prompt_nodes:
|
|
return None, None
|
|
|
|
# First pass: Score all nodes based on content
|
|
node_scores = []
|
|
for node in prompt_nodes:
|
|
prompt_text = _get_prompt_text(node)
|
|
# Skip empty or None text
|
|
if not prompt_text or not prompt_text.strip():
|
|
continue
|
|
|
|
prompt_lower = prompt_text.lower()
|
|
score = sum(1 for indicator in NEGATIVE_INDICATORS if indicator in prompt_lower)
|
|
node_scores.append((node, score, prompt_text))
|
|
|
|
if not node_scores:
|
|
return None, None
|
|
|
|
# Sort by score descending (highest score first for negative)
|
|
node_scores.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
positive_prompt = None
|
|
negative_prompt = None
|
|
|
|
# Highest score is likely negative (if score > 0)
|
|
if node_scores[0][1] > 0:
|
|
negative_prompt = node_scores[0][2]
|
|
# The node with the lowest score is likely the positive prompt
|
|
# Since it's sorted descending, pick the last one
|
|
positive_prompt = node_scores[-1][2]
|
|
|
|
# If we have more than 2 nodes and the last one is NOT the only zero-score node,
|
|
# we might need more logic, but for 2 nodes, this is correct.
|
|
# If positive and negative are the same (only 1 node with score > 0), clear positive.
|
|
if positive_prompt == negative_prompt:
|
|
positive_prompt = None
|
|
else:
|
|
# All scores are 0, use connection analysis
|
|
positive_prompt, negative_prompt = _classify_by_connections(
|
|
prompt_nodes, workflow_data, None, None
|
|
)
|
|
|
|
# Fallback: if we still can't determine, use first two nodes
|
|
if positive_prompt is None and negative_prompt is None and len(prompt_nodes) >= 2:
|
|
# Convention: assume first is positive, second is negative
|
|
positive_prompt = _get_prompt_text(prompt_nodes[0])
|
|
negative_prompt = _get_prompt_text(prompt_nodes[1])
|
|
elif positive_prompt is None and negative_prompt is not None:
|
|
# Find the other prompt
|
|
for _, _, text in node_scores:
|
|
if text != negative_prompt:
|
|
positive_prompt = text
|
|
break
|
|
elif negative_prompt is None and positive_prompt is not None:
|
|
# Find the other prompt
|
|
for _, _, text in node_scores:
|
|
if text != positive_prompt:
|
|
negative_prompt = text
|
|
break
|
|
|
|
return positive_prompt, negative_prompt
|
|
|
|
|
|
def _classify_by_connections(
|
|
prompt_nodes: List[Dict],
|
|
workflow_data: Dict,
|
|
existing_positive: Optional[str],
|
|
existing_negative: Optional[str]
|
|
) -> Tuple[Optional[str], Optional[str]]:
|
|
"""Try to classify prompts by analyzing node connections to samplers."""
|
|
positive = existing_positive
|
|
negative = existing_negative
|
|
|
|
links = workflow_data.get("links", [])
|
|
if not isinstance(links, list):
|
|
return positive, negative
|
|
|
|
nodes = workflow_data.get("nodes", [])
|
|
|
|
# Find sampler nodes
|
|
samplers = []
|
|
if isinstance(nodes, list):
|
|
for node in nodes:
|
|
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
|
|
samplers.append(node)
|
|
elif isinstance(nodes, dict):
|
|
for node_id, node in nodes.items():
|
|
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
|
|
node_copy = dict(node)
|
|
node_copy["id"] = node_id
|
|
samplers.append(node_copy)
|
|
|
|
if not samplers:
|
|
return positive, negative
|
|
|
|
# Trace connections from CLIP nodes to samplers
|
|
for link in links:
|
|
if len(link) < 4:
|
|
continue
|
|
|
|
from_node_id = link[0]
|
|
to_node_id = link[2]
|
|
to_slot = link[3]
|
|
|
|
# Find matching CLIP node and sampler
|
|
for clip_node in prompt_nodes:
|
|
clip_id = clip_node.get("id")
|
|
if clip_id == from_node_id:
|
|
for sampler in samplers:
|
|
if sampler.get("id") == to_node_id:
|
|
# Determine slot type
|
|
inputs = sampler.get("inputs", {})
|
|
if isinstance(inputs, dict):
|
|
input_keys = list(inputs.keys())
|
|
if to_slot < len(input_keys):
|
|
input_name = input_keys[to_slot].lower()
|
|
prompt_text = _get_prompt_text(clip_node)
|
|
|
|
if "positive" in input_name and positive is None:
|
|
positive = prompt_text
|
|
elif "negative" in input_name and negative is None:
|
|
negative = prompt_text
|
|
|
|
return positive, negative
|