621 lines
23 KiB
Python
621 lines
23 KiB
Python
"""
|
|
workflow_helpers.py
|
|
Description: Helper functions for workflow data processing, metadata enhancement, and visualization utilities.
|
|
Author: Eric Hiss (GitHub: EricRollei)
|
|
Contact: [eric@historic.camera, eric@rollei.us]
|
|
Version: 1.0.0
|
|
Date: [March 2025]
|
|
License: Dual License (Non-Commercial and Commercial Use)
|
|
Copyright (c) 2025 Eric Hiss. All rights reserved.
|
|
|
|
Dual License:
|
|
1. Non-Commercial Use: This software is licensed under the terms of the
|
|
Creative Commons Attribution-NonCommercial 4.0 International License.
|
|
To view a copy of this license, visit http://creativecommons.org/licenses/by-nc/4.0/
|
|
|
|
2. Commercial Use: For commercial use, a separate license is required.
|
|
Please contact Eric Hiss at [eric@historic.camera, eric@rollei.us] for licensing options.
|
|
|
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
|
|
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
|
PARTICULAR PURPOSE AND NONINFRINGEMENT.
|
|
|
|
Dependencies:
|
|
This code depends on several third-party libraries, each with its own license:
|
|
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import datetime
|
|
from typing import Dict, Any, List, Optional, Union, Tuple
|
|
|
|
def identify_workflow_source(workflow_data: Dict[str, Any]) -> str:
|
|
"""
|
|
Identify the source of a workflow (which AI system created it)
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
|
|
Returns:
|
|
str: Identified source ('comfyui', 'automatic1111', 'novelai', etc.)
|
|
"""
|
|
# Check for explicit type marker
|
|
if 'type' in workflow_data:
|
|
workflow_type = workflow_data.get('type', '')
|
|
if 'comfy' in workflow_type.lower():
|
|
return 'comfyui'
|
|
if 'a1111' in workflow_type.lower() or 'automatic' in workflow_type.lower():
|
|
return 'automatic1111'
|
|
if 'stable' in workflow_type.lower() and 'diffusion' in workflow_type.lower():
|
|
return 'stable_diffusion_webui'
|
|
if 'novelai' in workflow_type.lower():
|
|
return 'novelai'
|
|
return workflow_type
|
|
|
|
# Look for ComfyUI specific structure
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict) and 'nodes' in workflow_data.get('prompt', {}):
|
|
return 'comfyui'
|
|
|
|
# Look for A1111 specific structure
|
|
if all(k in workflow_data for k in ['prompt', 'negative_prompt', 'sampler']):
|
|
return 'automatic1111'
|
|
|
|
# Look for NovelAI specific markers
|
|
if 'novelai' in str(workflow_data).lower() or 'nai' in str(workflow_data).lower():
|
|
return 'novelai'
|
|
|
|
return 'unknown'
|
|
|
|
def extract_generation_parameters(workflow_data: Dict[str, Any], source: str) -> Dict[str, Any]:
|
|
"""
|
|
Extract generation parameters from workflow based on its source
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
source: Identified source ('comfyui', 'automatic1111', etc.)
|
|
|
|
Returns:
|
|
dict: Extracted generation parameters
|
|
"""
|
|
generation_params = {}
|
|
|
|
if source == 'comfyui':
|
|
# Extract from ComfyUI workflow
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict):
|
|
nodes = workflow_data['prompt'].get('nodes', {})
|
|
|
|
# Find checkpoint nodes
|
|
for node_id, node in nodes.items():
|
|
node_type = node.get('class_type', '')
|
|
inputs = node.get('inputs', {})
|
|
|
|
# Extract model
|
|
if 'Checkpoint' in node_type and 'ckpt_name' in inputs:
|
|
generation_params['model'] = inputs['ckpt_name']
|
|
|
|
# Extract sampler settings
|
|
elif 'KSampler' in node_type:
|
|
generation_params['sampler'] = inputs.get('sampler_name')
|
|
generation_params['scheduler'] = inputs.get('scheduler')
|
|
generation_params['steps'] = inputs.get('steps')
|
|
generation_params['cfg_scale'] = inputs.get('cfg')
|
|
generation_params['seed'] = inputs.get('seed')
|
|
|
|
# Extract VAE
|
|
elif 'VAE' in node_type and 'vae_name' in inputs:
|
|
generation_params['vae'] = inputs['vae_name']
|
|
|
|
# Extract prompt
|
|
elif 'CLIPTextEncode' in node_type and 'text' in inputs:
|
|
if node.get('is_negative', False):
|
|
generation_params['negative_prompt'] = inputs['text']
|
|
else:
|
|
generation_params['prompt'] = inputs['text']
|
|
|
|
elif source in ['automatic1111', 'stable_diffusion_webui']:
|
|
# Direct copy of parameters from A1111/SD-WebUI
|
|
for key in ['model', 'prompt', 'negative_prompt', 'sampler',
|
|
'steps', 'cfg_scale', 'seed', 'vae']:
|
|
if key in workflow_data:
|
|
generation_params[key] = workflow_data[key]
|
|
|
|
# Add timestamp
|
|
generation_params['timestamp'] = datetime.datetime.now().isoformat()
|
|
|
|
return generation_params
|
|
|
|
def calculate_workflow_complexity(workflow_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Calculate complexity metrics for a workflow
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
|
|
Returns:
|
|
dict: Complexity metrics including node count, connection count,
|
|
complexity score, etc.
|
|
"""
|
|
stats = {
|
|
'node_count': 0,
|
|
'connection_count': 0,
|
|
'unique_node_types': 0,
|
|
'complexity_score': 0,
|
|
'node_types': {}
|
|
}
|
|
|
|
# Check for ComfyUI structure
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict):
|
|
nodes = workflow_data['prompt'].get('nodes', {})
|
|
links = workflow_data['prompt'].get('links', [])
|
|
|
|
# Count nodes and node types
|
|
stats['node_count'] = len(nodes)
|
|
node_types = {}
|
|
|
|
for node_id, node in nodes.items():
|
|
node_type = node.get('class_type', 'Unknown')
|
|
node_types[node_type] = node_types.get(node_type, 0) + 1
|
|
|
|
stats['node_types'] = node_types
|
|
stats['unique_node_types'] = len(node_types)
|
|
|
|
# Count connections
|
|
stats['connection_count'] = len(links)
|
|
|
|
# Calculate complexity score
|
|
# More nodes, more connections, and more unique node types = more complex
|
|
stats['complexity_score'] = (
|
|
stats['node_count'] * 2 +
|
|
stats['connection_count'] * 1.5 +
|
|
stats['unique_node_types'] * 3
|
|
)
|
|
|
|
return stats
|
|
|
|
def extract_embedded_workflow(filepath: str) -> Dict[str, Any]:
|
|
"""
|
|
Extract workflow data from embedded metadata in an image file
|
|
|
|
Args:
|
|
filepath: Path to the image file
|
|
|
|
Returns:
|
|
dict: Extracted workflow data or empty dict if none found
|
|
"""
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(filepath) as img:
|
|
# Check for workflow data in PNG text chunks
|
|
if 'prompt' in img.info:
|
|
try:
|
|
return json.loads(img.info['prompt'])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Check for workflow in extra PNG info
|
|
if 'workflow' in img.info:
|
|
try:
|
|
return json.loads(img.info['workflow'])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Check for parameters in A1111 format
|
|
if 'parameters' in img.info:
|
|
try:
|
|
params = img.info['parameters']
|
|
# Parse A1111 parameters string into a structured format
|
|
result = {'type': 'automatic1111'}
|
|
|
|
# Split on the first 'Negative prompt:'
|
|
parts = params.split('Negative prompt:', 1)
|
|
if len(parts) == 2:
|
|
result['prompt'] = parts[0].strip()
|
|
# The rest contains negative prompt and settings
|
|
neg_and_settings = parts[1].split('Steps:', 1)
|
|
if len(neg_and_settings) == 2:
|
|
result['negative_prompt'] = neg_and_settings[0].strip()
|
|
# Parse the technical settings
|
|
settings_text = 'Steps:' + neg_and_settings[1]
|
|
settings_pattern = r'([a-zA-Z\s]+):\s*([^,]+)(?:,|$)'
|
|
for setting_match in re.finditer(settings_pattern, settings_text):
|
|
key = setting_match.group(1).strip().lower().replace(' ', '_')
|
|
value = setting_match.group(2).strip()
|
|
result[key] = value
|
|
else:
|
|
result['negative_prompt'] = neg_and_settings[0].strip()
|
|
else:
|
|
result['prompt'] = params
|
|
|
|
return result
|
|
except Exception:
|
|
pass
|
|
|
|
return {}
|
|
except Exception as e:
|
|
print(f"Error extracting embedded workflow: {str(e)}")
|
|
return {}
|
|
|
|
def extract_workflow_node_types(workflow_data: Dict[str, Any]) -> Dict[str, int]:
|
|
"""
|
|
Extract and count all node types used in a workflow
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
|
|
Returns:
|
|
dict: Mapping of node type names to their count in the workflow
|
|
"""
|
|
node_types = {}
|
|
|
|
# Handle ComfyUI workflow structure
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict):
|
|
nodes = workflow_data['prompt'].get('nodes', {})
|
|
|
|
for node_id, node in nodes.items():
|
|
if isinstance(node, dict) and 'class_type' in node:
|
|
node_type = node['class_type']
|
|
node_types[node_type] = node_types.get(node_type, 0) + 1
|
|
|
|
# Direct node structure
|
|
elif isinstance(workflow_data, dict):
|
|
for key, value in workflow_data.items():
|
|
if isinstance(value, dict) and 'class_type' in value:
|
|
node_type = value['class_type']
|
|
node_types[node_type] = node_types.get(node_type, 0) + 1
|
|
|
|
return node_types
|
|
|
|
def extract_model_names(workflow_data: Dict[str, Any]) -> List[str]:
|
|
"""
|
|
Extract all model names (checkpoints, VAEs, LoRAs, etc.) from a workflow
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
|
|
Returns:
|
|
list: List of model names found in the workflow
|
|
"""
|
|
models = []
|
|
|
|
# Check for ComfyUI structure
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict):
|
|
nodes = workflow_data['prompt'].get('nodes', {})
|
|
|
|
for node_id, node in nodes.items():
|
|
if not isinstance(node, dict) or 'class_type' not in node:
|
|
continue
|
|
|
|
class_type = node['class_type']
|
|
inputs = node.get('inputs', {})
|
|
|
|
# Checkpoint loader
|
|
if 'CheckpointLoader' in class_type and 'ckpt_name' in inputs:
|
|
models.append(inputs['ckpt_name'])
|
|
|
|
# VAE loader
|
|
elif 'VAELoader' in class_type and 'vae_name' in inputs:
|
|
models.append(inputs['vae_name'])
|
|
|
|
# LoRA loader
|
|
elif 'LoRA' in class_type and 'lora_name' in inputs:
|
|
models.append(inputs['lora_name'])
|
|
|
|
# ControlNet model
|
|
elif 'ControlNet' in class_type and 'control_net_name' in inputs:
|
|
models.append(inputs['control_net_name'])
|
|
|
|
# Remove duplicates while preserving order
|
|
seen = set()
|
|
unique_models = [m for m in models if not (m in seen or seen.add(m))]
|
|
|
|
return unique_models
|
|
|
|
def extract_prompts(workflow_data: Dict[str, Any]) -> Dict[str, List[str]]:
|
|
"""
|
|
Extract positive and negative prompts from a workflow
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
|
|
Returns:
|
|
dict: Dictionary with 'positive' and 'negative' prompt lists
|
|
"""
|
|
result = {
|
|
'positive': [],
|
|
'negative': []
|
|
}
|
|
|
|
# Check for ComfyUI structure
|
|
if 'prompt' in workflow_data and isinstance(workflow_data['prompt'], dict):
|
|
nodes = workflow_data['prompt'].get('nodes', {})
|
|
|
|
for node_id, node in nodes.items():
|
|
if not isinstance(node, dict) or 'class_type' not in node:
|
|
continue
|
|
|
|
# Look for CLIP text encode nodes
|
|
if 'CLIPTextEncode' in node['class_type']:
|
|
inputs = node.get('inputs', {})
|
|
if 'text' in inputs:
|
|
text = inputs['text']
|
|
|
|
# Determine if positive or negative
|
|
is_negative = False
|
|
|
|
# Check node title for indication
|
|
title = node.get('title', '').lower()
|
|
if 'negative' in title or 'neg' in title:
|
|
is_negative = True
|
|
|
|
# Check connections for indication
|
|
for link_id, link in enumerate(workflow_data['prompt'].get('links', [])):
|
|
if link[0] == node_id: # Node is source
|
|
# Check if target is negative input
|
|
target_id = link[2]
|
|
target_slot = link[3]
|
|
target = nodes.get(target_id, {})
|
|
if target.get('inputs', {}).get('neg_conditioning') == target_slot:
|
|
is_negative = True
|
|
break
|
|
|
|
# Add to appropriate list
|
|
if is_negative:
|
|
result['negative'].append(text)
|
|
else:
|
|
result['positive'].append(text)
|
|
|
|
# Check for direct prompt fields (A1111 style)
|
|
elif isinstance(workflow_data, dict):
|
|
if 'prompt' in workflow_data:
|
|
result['positive'].append(workflow_data['prompt'])
|
|
if 'negative_prompt' in workflow_data:
|
|
result['negative'].append(workflow_data['negative_prompt'])
|
|
|
|
return result
|
|
|
|
def format_workflow_summary(workflow_data: Dict[str, Any], format_type: str = 'text') -> str:
|
|
"""
|
|
Create a user-friendly summary of a workflow
|
|
|
|
Args:
|
|
workflow_data: Raw workflow data dictionary
|
|
format_type: Output format ('text', 'html', 'markdown')
|
|
|
|
Returns:
|
|
str: Formatted workflow summary
|
|
"""
|
|
# Identify workflow source
|
|
source = identify_workflow_source(workflow_data)
|
|
|
|
# Extract key components
|
|
models = extract_model_names(workflow_data)
|
|
prompts = extract_prompts(workflow_data)
|
|
node_types = extract_workflow_node_types(workflow_data)
|
|
complexity = calculate_workflow_complexity(workflow_data)
|
|
parameters = extract_generation_parameters(workflow_data, source)
|
|
|
|
# Format based on output type
|
|
if format_type == 'text':
|
|
return _format_workflow_summary_as_text(
|
|
source, models, prompts, node_types, complexity, parameters)
|
|
elif format_type == 'html':
|
|
return _format_workflow_summary_as_html(
|
|
source, models, prompts, node_types, complexity, parameters)
|
|
elif format_type == 'markdown':
|
|
return _format_workflow_summary_as_markdown(
|
|
source, models, prompts, node_types, complexity, parameters)
|
|
else:
|
|
return f"Unknown format type: {format_type}"
|
|
|
|
def _format_workflow_summary_as_text(
|
|
source: str,
|
|
models: List[str],
|
|
prompts: Dict[str, List[str]],
|
|
node_types: Dict[str, int],
|
|
complexity: Dict[str, Any],
|
|
parameters: Dict[str, Any]) -> str:
|
|
"""Format workflow summary as plain text"""
|
|
lines = [f"Workflow Source: {source.upper()}", ""]
|
|
|
|
# Add models
|
|
if models:
|
|
lines.append("Models:")
|
|
for model in models:
|
|
lines.append(f"- {model}")
|
|
lines.append("")
|
|
|
|
# Add generation parameters
|
|
if parameters:
|
|
lines.append("Generation Parameters:")
|
|
for key, value in parameters.items():
|
|
if key not in ('prompt', 'negative_prompt'):
|
|
lines.append(f"- {key}: {value}")
|
|
lines.append("")
|
|
|
|
# Add prompts
|
|
if prompts['positive']:
|
|
lines.append("Positive Prompt:")
|
|
for prompt in prompts['positive']:
|
|
lines.append(prompt)
|
|
lines.append("")
|
|
|
|
if prompts['negative']:
|
|
lines.append("Negative Prompt:")
|
|
for prompt in prompts['negative']:
|
|
lines.append(prompt)
|
|
lines.append("")
|
|
|
|
# Add complexity info
|
|
if complexity:
|
|
lines.append("Workflow Complexity:")
|
|
if 'node_count' in complexity:
|
|
lines.append(f"- Nodes: {complexity['node_count']}")
|
|
if 'connection_count' in complexity:
|
|
lines.append(f"- Connections: {complexity['connection_count']}")
|
|
if 'complexity_score' in complexity:
|
|
lines.append(f"- Complexity Score: {complexity['complexity_score']}")
|
|
if 'unique_node_types' in complexity:
|
|
lines.append(f"- Unique Node Types: {complexity['unique_node_types']}")
|
|
lines.append("")
|
|
|
|
# Add node type summary (top 10)
|
|
if node_types:
|
|
lines.append("Top Node Types:")
|
|
sorted_types = sorted(node_types.items(), key=lambda x: x[1], reverse=True)
|
|
for node_type, count in sorted_types[:10]:
|
|
lines.append(f"- {node_type}: {count}")
|
|
|
|
return "\n".join(lines)
|
|
|
|
def _format_workflow_summary_as_markdown(
|
|
source: str,
|
|
models: List[str],
|
|
prompts: Dict[str, List[str]],
|
|
node_types: Dict[str, int],
|
|
complexity: Dict[str, Any],
|
|
parameters: Dict[str, Any]) -> str:
|
|
"""Format workflow summary as markdown"""
|
|
lines = [f"# Workflow Analysis ({source.upper()})", ""]
|
|
|
|
# Add models
|
|
if models:
|
|
lines.append("## Models")
|
|
for model in models:
|
|
lines.append(f"- {model}")
|
|
lines.append("")
|
|
|
|
# Add generation parameters
|
|
if parameters:
|
|
lines.append("## Generation Parameters")
|
|
for key, value in parameters.items():
|
|
if key not in ('prompt', 'negative_prompt'):
|
|
lines.append(f"- **{key}:** {value}")
|
|
lines.append("")
|
|
|
|
# Add prompts
|
|
if prompts['positive']:
|
|
lines.append("## Positive Prompt")
|
|
for prompt in prompts['positive']:
|
|
lines.append(f"```\n{prompt}\n```")
|
|
lines.append("")
|
|
|
|
if prompts['negative']:
|
|
lines.append("## Negative Prompt")
|
|
for prompt in prompts['negative']:
|
|
lines.append(f"```\n{prompt}\n```")
|
|
lines.append("")
|
|
|
|
# Add complexity info
|
|
if complexity:
|
|
lines.append("## Workflow Complexity")
|
|
if 'node_count' in complexity:
|
|
lines.append(f"- **Nodes:** {complexity['node_count']}")
|
|
if 'connection_count' in complexity:
|
|
lines.append(f"- **Connections:** {complexity['connection_count']}")
|
|
if 'complexity_score' in complexity:
|
|
lines.append(f"- **Complexity Score:** {complexity['complexity_score']}")
|
|
if 'unique_node_types' in complexity:
|
|
lines.append(f"- **Unique Node Types:** {complexity['unique_node_types']}")
|
|
lines.append("")
|
|
|
|
# Add node type summary (top 10)
|
|
if node_types:
|
|
lines.append("## Top Node Types")
|
|
sorted_types = sorted(node_types.items(), key=lambda x: x[1], reverse=True)
|
|
for node_type, count in sorted_types[:10]:
|
|
lines.append(f"- **{node_type}:** {count}")
|
|
|
|
return "\n".join(lines)
|
|
|
|
def _format_workflow_summary_as_html(
|
|
source: str,
|
|
models: List[str],
|
|
prompts: Dict[str, List[str]],
|
|
node_types: Dict[str, int],
|
|
complexity: Dict[str, Any],
|
|
parameters: Dict[str, Any]) -> str:
|
|
"""Format workflow summary as HTML"""
|
|
html = [
|
|
"<html><head><title>Workflow Analysis</title>",
|
|
"<style>",
|
|
"body { font-family: Arial, sans-serif; margin: 20px; }",
|
|
"h1, h2 { color: #333366; }",
|
|
"h1 { border-bottom: 1px solid #cccccc; padding-bottom: 5px; }",
|
|
"h2 { margin-top: 20px; }",
|
|
"ul { margin-top: 5px; }",
|
|
"pre { background-color: #f8f8f8; padding: 10px; border: 1px solid #ddd; }",
|
|
".section { margin-bottom: 20px; }",
|
|
"</style>",
|
|
"</head><body>",
|
|
f"<h1>Workflow Analysis ({source.upper()})</h1>"
|
|
]
|
|
|
|
# Add models
|
|
if models:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Models</h2>")
|
|
html.append("<ul>")
|
|
for model in models:
|
|
html.append(f"<li>{model}</li>")
|
|
html.append("</ul>")
|
|
html.append("</div>")
|
|
|
|
# Add generation parameters
|
|
if parameters:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Generation Parameters</h2>")
|
|
html.append("<ul>")
|
|
for key, value in parameters.items():
|
|
if key not in ('prompt', 'negative_prompt'):
|
|
html.append(f"<li><strong>{key}:</strong> {value}</li>")
|
|
html.append("</ul>")
|
|
html.append("</div>")
|
|
|
|
# Add prompts
|
|
if prompts['positive']:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Positive Prompt</h2>")
|
|
for prompt in prompts['positive']:
|
|
html.append(f"<pre>{prompt}</pre>")
|
|
html.append("</div>")
|
|
|
|
if prompts['negative']:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Negative Prompt</h2>")
|
|
for prompt in prompts['negative']:
|
|
html.append(f"<pre>{prompt}</pre>")
|
|
html.append("</div>")
|
|
|
|
# Add complexity info
|
|
if complexity:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Workflow Complexity</h2>")
|
|
html.append("<ul>")
|
|
if 'node_count' in complexity:
|
|
html.append(f"<li><strong>Nodes:</strong> {complexity['node_count']}</li>")
|
|
if 'connection_count' in complexity:
|
|
html.append(f"<li><strong>Connections:</strong> {complexity['connection_count']}</li>")
|
|
if 'complexity_score' in complexity:
|
|
html.append(f"<li><strong>Complexity Score:</strong> {complexity['complexity_score']}</li>")
|
|
if 'unique_node_types' in complexity:
|
|
html.append(f"<li><strong>Unique Node Types:</strong> {complexity['unique_node_types']}</li>")
|
|
html.append("</ul>")
|
|
html.append("</div>")
|
|
|
|
# Add node type summary (top 10)
|
|
if node_types:
|
|
html.append("<div class='section'>")
|
|
html.append("<h2>Top Node Types</h2>")
|
|
html.append("<ul>")
|
|
sorted_types = sorted(node_types.items(), key=lambda x: x[1], reverse=True)
|
|
for node_type, count in sorted_types[:10]:
|
|
html.append(f"<li><strong>{node_type}:</strong> {count}</li>")
|
|
html.append("</ul>")
|
|
html.append("</div>")
|
|
|
|
html.append("</body></html>")
|
|
return "\n".join(html)
|