PXL8D
3c7b5658f0
- Updated all type hints to use built-in types (dict, list, tuple) instead of typing imports
- Removed Dict, List, Tuple from typing imports since they're not needed in Python 3.10+
- Fixed NameError: name 'Dict' is not defined in specialized_nodes.py
- Modernized type annotations throughout codebase for better compatibility
- Resolves ComfyUI import error when loading MetaMan nodes
Files updated:
- specialized_nodes.py: Fixed all Dict/List/Tuple type hints
- metaman.py: Updated type annotations to modern syntax
- metadata_parser.py: Modernized type hints for consistency
Tested: ✅ Nodes now load successfully in ComfyUI without import errors
569 lines
22 KiB
Python
569 lines
22 KiB
Python
"""
|
|
MetaMan Workflow Saver Node
|
|
Specialized node for saving and embedding workflows in images
|
|
"""
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
# Fallback for testing outside ComfyUI environment
|
|
torch = None
|
|
import json
|
|
import os
|
|
from typing import Optional, Any, Union
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from datetime import datetime
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
# Fallback for testing outside ComfyUI environment
|
|
class folder_paths:
|
|
@staticmethod
|
|
def get_output_directory():
|
|
return os.getcwd()
|
|
from .metadata_parser import metadata_parser
|
|
|
|
|
|
class MetaManWorkflowSaver:
|
|
"""
|
|
Dedicated node for saving workflows with various output options
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"save_format": ([
|
|
"embed_in_image", "json_file", "both", "workflow_only", "prompt_only"
|
|
], {"default": "embed_in_image"}),
|
|
"filename_prefix": ("STRING", {"default": "MetaMan_workflow"}),
|
|
},
|
|
"optional": {
|
|
"workflow_data": ("STRING", {"default": ""}),
|
|
"include_metadata": ("BOOLEAN", {"default": True}),
|
|
"compress_data": ("BOOLEAN", {"default": True}),
|
|
"preserve_original": ("BOOLEAN", {"default": True}),
|
|
"output_directory": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("image", "workflow_info", "file_path")
|
|
FUNCTION = "save_workflow"
|
|
CATEGORY = "MetaMan"
|
|
DESCRIPTION = "Save workflows as files or embed in images"
|
|
|
|
def save_workflow(self, image, save_format, filename_prefix,
|
|
workflow_data="", include_metadata=True, compress_data=True,
|
|
preserve_original=True, output_directory=""):
|
|
"""
|
|
Save workflow data in various formats
|
|
"""
|
|
try:
|
|
# Convert tensor to PIL Image
|
|
if isinstance(image, torch.Tensor):
|
|
img_tensor = image[0]
|
|
img_array = (img_tensor.cpu().numpy() * 255).astype('uint8')
|
|
pil_image = Image.fromarray(img_array)
|
|
else:
|
|
pil_image = image
|
|
|
|
# Determine output directory
|
|
if not output_directory:
|
|
output_directory = folder_paths.get_output_directory()
|
|
|
|
# Generate filename
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
base_filename = f"{filename_prefix}_{timestamp}"
|
|
|
|
workflow_info = ""
|
|
file_path = ""
|
|
|
|
# Extract or prepare workflow data
|
|
if not workflow_data:
|
|
workflow_data = self._extract_workflow_from_image(pil_image)
|
|
|
|
if save_format == "embed_in_image":
|
|
result_image, info = self._embed_workflow_in_image(
|
|
pil_image, workflow_data, include_metadata, compress_data, preserve_original
|
|
)
|
|
# Save the image with embedded workflow
|
|
image_path = os.path.join(output_directory, f"{base_filename}.png")
|
|
result_image.save(image_path, "PNG")
|
|
|
|
workflow_info = info
|
|
file_path = image_path
|
|
output_image = result_image
|
|
|
|
elif save_format == "json_file":
|
|
json_path = os.path.join(output_directory, f"{base_filename}.json")
|
|
workflow_info = self._save_workflow_json(workflow_data, json_path, include_metadata)
|
|
file_path = json_path
|
|
output_image = pil_image
|
|
|
|
elif save_format == "both":
|
|
# Embed in image
|
|
result_image, embed_info = self._embed_workflow_in_image(
|
|
pil_image, workflow_data, include_metadata, compress_data, preserve_original
|
|
)
|
|
image_path = os.path.join(output_directory, f"{base_filename}.png")
|
|
result_image.save(image_path, "PNG")
|
|
|
|
# Save JSON file
|
|
json_path = os.path.join(output_directory, f"{base_filename}.json")
|
|
json_info = self._save_workflow_json(workflow_data, json_path, include_metadata)
|
|
|
|
workflow_info = f"Image: {embed_info}\nJSON: {json_info}"
|
|
file_path = f"Image: {image_path}, JSON: {json_path}"
|
|
output_image = result_image
|
|
|
|
elif save_format == "workflow_only":
|
|
workflow_path = os.path.join(output_directory, f"{base_filename}_workflow.json")
|
|
workflow_info = self._save_workflow_only(workflow_data, workflow_path)
|
|
file_path = workflow_path
|
|
output_image = pil_image
|
|
|
|
elif save_format == "prompt_only":
|
|
prompt_path = os.path.join(output_directory, f"{base_filename}_prompt.json")
|
|
workflow_info = self._save_prompt_only(workflow_data, prompt_path)
|
|
file_path = prompt_path
|
|
output_image = pil_image
|
|
|
|
# Convert back to tensor if needed
|
|
if isinstance(image, torch.Tensor):
|
|
if isinstance(output_image, Image.Image):
|
|
img_array = torch.from_numpy(output_image.__array__()).float() / 255.0
|
|
output_tensor = img_array.unsqueeze(0)
|
|
return (output_tensor, workflow_info, file_path)
|
|
|
|
return (output_image, workflow_info, file_path)
|
|
|
|
except Exception as e:
|
|
error_msg = f"MetaMan Workflow Saver Error: {str(e)}"
|
|
return (image, error_msg, "")
|
|
|
|
def _extract_workflow_from_image(self, image: Image.Image) -> dict:
|
|
"""Extract existing workflow data from image"""
|
|
extracted = metadata_parser.extract_all_metadata(image)
|
|
|
|
workflow_data = {}
|
|
|
|
# Get ComfyUI workflow if available
|
|
if 'comfyui_workflow' in extracted['parsed_metadata']:
|
|
workflow_data['workflow'] = extracted['parsed_metadata']['comfyui_workflow']
|
|
|
|
# Get ComfyUI prompt if available
|
|
if 'comfyui_prompt' in extracted['parsed_metadata']:
|
|
workflow_data['prompt'] = extracted['parsed_metadata']['comfyui_prompt']
|
|
|
|
# Include all parsed metadata
|
|
workflow_data['metadata'] = extracted['parsed_metadata']
|
|
workflow_data['source_service'] = extracted['source_detected']
|
|
workflow_data['extracted_at'] = datetime.now().isoformat()
|
|
|
|
return workflow_data
|
|
|
|
def _embed_workflow_in_image(self, image: Image.Image, workflow_data: dict,
|
|
include_metadata: bool, compress_data: bool,
|
|
preserve_original: bool) -> tuple[Image.Image, str]:
|
|
"""Embed workflow data into image PNG chunks"""
|
|
|
|
# Create new image or copy original
|
|
if preserve_original and hasattr(image, 'copy'):
|
|
result_image = image.copy()
|
|
else:
|
|
result_image = image
|
|
|
|
# Prepare PNG info
|
|
png_info = PngInfo()
|
|
|
|
# Copy existing metadata if preserving original
|
|
if preserve_original and hasattr(image, 'text'):
|
|
for key, value in image.text.items():
|
|
png_info.add_text(key, value)
|
|
|
|
chunks_added = []
|
|
|
|
# Add workflow data
|
|
if isinstance(workflow_data, dict):
|
|
if 'workflow' in workflow_data:
|
|
workflow_json = json.dumps(workflow_data['workflow'], separators=(',', ':'))
|
|
png_info.add_text("workflow", workflow_json, zip=compress_data)
|
|
chunks_added.append("workflow")
|
|
|
|
if 'prompt' in workflow_data:
|
|
prompt_json = json.dumps(workflow_data['prompt'], separators=(',', ':'))
|
|
png_info.add_text("prompt", prompt_json, zip=compress_data)
|
|
chunks_added.append("prompt")
|
|
|
|
# Add universal metadata if requested
|
|
if include_metadata and 'metadata' in workflow_data:
|
|
# Create MetaMan universal chunk
|
|
universal_meta = {
|
|
'schema_version': '1.0.0',
|
|
'created_at': datetime.now().isoformat(),
|
|
'metaman_version': '1.0.0',
|
|
'source_service': workflow_data.get('source_service', 'unknown'),
|
|
'metadata': workflow_data['metadata']
|
|
}
|
|
meta_json = json.dumps(universal_meta, separators=(',', ':'))
|
|
png_info.add_text("meta", meta_json, zip=compress_data)
|
|
chunks_added.append("meta")
|
|
|
|
# Create new image with metadata
|
|
if chunks_added:
|
|
# Save to temporary location to embed metadata
|
|
import tempfile
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
|
|
result_image.save(tmp_file, "PNG", pnginfo=png_info)
|
|
tmp_file.flush()
|
|
|
|
# Reload image with embedded metadata
|
|
result_image = Image.open(tmp_file.name)
|
|
result_image.load() # Ensure image is fully loaded
|
|
|
|
# Clean up temp file
|
|
os.unlink(tmp_file.name)
|
|
|
|
info_message = f"Embedded chunks: {', '.join(chunks_added)}"
|
|
|
|
return result_image, info_message
|
|
|
|
def _save_workflow_json(self, workflow_data: dict, file_path: str, include_metadata: bool) -> str:
|
|
"""Save complete workflow data as JSON file"""
|
|
|
|
output_data = {
|
|
'metaman_export': {
|
|
'version': '1.0.0',
|
|
'exported_at': datetime.now().isoformat(),
|
|
'export_type': 'complete_workflow'
|
|
}
|
|
}
|
|
|
|
if isinstance(workflow_data, dict):
|
|
output_data.update(workflow_data)
|
|
else:
|
|
output_data['workflow_data'] = workflow_data
|
|
|
|
if not include_metadata:
|
|
output_data.pop('metadata', None)
|
|
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
json.dump(output_data, f, indent=2, ensure_ascii=False)
|
|
|
|
return f"Saved complete workflow to {os.path.basename(file_path)}"
|
|
|
|
def _save_workflow_only(self, workflow_data: dict, file_path: str) -> str:
|
|
"""Save only the workflow graph data"""
|
|
|
|
if isinstance(workflow_data, dict) and 'workflow' in workflow_data:
|
|
workflow_only = workflow_data['workflow']
|
|
else:
|
|
workflow_only = workflow_data
|
|
|
|
output_data = {
|
|
'metaman_export': {
|
|
'version': '1.0.0',
|
|
'exported_at': datetime.now().isoformat(),
|
|
'export_type': 'workflow_only'
|
|
},
|
|
'workflow': workflow_only
|
|
}
|
|
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
json.dump(output_data, f, indent=2, ensure_ascii=False)
|
|
|
|
return f"Saved workflow graph to {os.path.basename(file_path)}"
|
|
|
|
def _save_prompt_only(self, workflow_data: dict, file_path: str) -> str:
|
|
"""Save only the prompt execution data"""
|
|
|
|
if isinstance(workflow_data, dict) and 'prompt' in workflow_data:
|
|
prompt_only = workflow_data['prompt']
|
|
else:
|
|
prompt_only = {}
|
|
|
|
output_data = {
|
|
'metaman_export': {
|
|
'version': '1.0.0',
|
|
'exported_at': datetime.now().isoformat(),
|
|
'export_type': 'prompt_only'
|
|
},
|
|
'prompt': prompt_only
|
|
}
|
|
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
json.dump(output_data, f, indent=2, ensure_ascii=False)
|
|
|
|
return f"Saved prompt data to {os.path.basename(file_path)}"
|
|
|
|
|
|
class MetaManDependencyResolver:
|
|
"""
|
|
Specialized node for resolving model dependencies and finding download sources
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"search_platforms": ([
|
|
"all", "civitai_only", "huggingface_only", "civitai_then_hf"
|
|
], {"default": "all"}),
|
|
"output_format": ([
|
|
"detailed_json", "download_urls", "summary_text", "dependency_list"
|
|
], {"default": "detailed_json"}),
|
|
},
|
|
"optional": {
|
|
"include_confidence": ("BOOLEAN", {"default": True}),
|
|
"min_confidence": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
"max_results_per_model": ("INT", {"default": 3, "min": 1, "max": 10}),
|
|
"cache_results": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING")
|
|
RETURN_NAMES = ("image", "dependencies", "summary")
|
|
FUNCTION = "resolve_dependencies"
|
|
CATEGORY = "MetaMan"
|
|
DESCRIPTION = "Resolve model dependencies and find download sources"
|
|
|
|
def resolve_dependencies(self, image, search_platforms, output_format,
|
|
include_confidence=True, min_confidence=0.3,
|
|
max_results_per_model=3, cache_results=True):
|
|
"""
|
|
Resolve model dependencies from image metadata
|
|
"""
|
|
try:
|
|
from .api_integration import model_resolver
|
|
|
|
# Convert tensor to PIL Image
|
|
if isinstance(image, torch.Tensor):
|
|
img_tensor = image[0]
|
|
img_array = (img_tensor.cpu().numpy() * 255).astype('uint8')
|
|
pil_image = Image.fromarray(img_array)
|
|
else:
|
|
pil_image = image
|
|
|
|
# Extract metadata
|
|
extracted = metadata_parser.extract_all_metadata(pil_image)
|
|
metadata = extracted['parsed_metadata']
|
|
|
|
# Find dependencies
|
|
dependencies = self._find_all_dependencies(metadata)
|
|
|
|
# Resolve each dependency
|
|
resolved_dependencies = []
|
|
|
|
for dep in dependencies:
|
|
if search_platforms in ["all", "civitai_only", "civitai_then_hf"]:
|
|
sources = model_resolver.find_model_sources(
|
|
dep['name'],
|
|
dep.get('hash', ''),
|
|
dep['type']
|
|
)
|
|
|
|
# Filter by confidence and limit results
|
|
filtered_sources = [
|
|
s for s in sources
|
|
if s.get('confidence', 0) >= min_confidence
|
|
][:max_results_per_model]
|
|
|
|
dep['sources'] = filtered_sources
|
|
dep['search_status'] = 'found' if filtered_sources else 'not_found'
|
|
else:
|
|
dep['sources'] = []
|
|
dep['search_status'] = 'skipped'
|
|
|
|
resolved_dependencies.append(dep)
|
|
|
|
# Format output
|
|
if output_format == "detailed_json":
|
|
output = self._format_detailed_json(resolved_dependencies, include_confidence)
|
|
elif output_format == "download_urls":
|
|
output = self._format_download_urls(resolved_dependencies)
|
|
elif output_format == "summary_text":
|
|
output = self._format_summary_text(resolved_dependencies)
|
|
elif output_format == "dependency_list":
|
|
output = self._format_dependency_list(resolved_dependencies)
|
|
else:
|
|
output = json.dumps(resolved_dependencies, indent=2)
|
|
|
|
# Generate summary
|
|
summary = self._generate_summary(resolved_dependencies)
|
|
|
|
return (image, output, summary)
|
|
|
|
except Exception as e:
|
|
error_msg = f"MetaMan Dependency Resolver Error: {str(e)}"
|
|
return (image, error_msg, "")
|
|
|
|
def _find_all_dependencies(self, metadata: dict) -> list[dict]:
|
|
"""Find all model dependencies in metadata"""
|
|
dependencies = []
|
|
|
|
# Main model
|
|
if 'model_name' in metadata and metadata['model_name']:
|
|
dep = {
|
|
'type': 'checkpoint',
|
|
'name': metadata['model_name'],
|
|
'hash': metadata.get('model_hash', ''),
|
|
'required': True,
|
|
'weight': None
|
|
}
|
|
dependencies.append(dep)
|
|
|
|
# LoRAs
|
|
if 'loras' in metadata and isinstance(metadata['loras'], list):
|
|
for lora in metadata['loras']:
|
|
if isinstance(lora, dict) and lora.get('name'):
|
|
dep = {
|
|
'type': 'lora',
|
|
'name': lora['name'],
|
|
'hash': lora.get('hash', ''),
|
|
'weight': lora.get('weight', 1.0),
|
|
'required': True
|
|
}
|
|
dependencies.append(dep)
|
|
|
|
# Embeddings
|
|
if 'embeddings' in metadata and isinstance(metadata['embeddings'], list):
|
|
for embedding in metadata['embeddings']:
|
|
if isinstance(embedding, dict) and embedding.get('name'):
|
|
dep = {
|
|
'type': 'embedding',
|
|
'name': embedding['name'],
|
|
'hash': embedding.get('hash', ''),
|
|
'weight': None,
|
|
'required': True
|
|
}
|
|
dependencies.append(dep)
|
|
|
|
# VAE
|
|
if 'vae' in metadata and metadata['vae']:
|
|
dep = {
|
|
'type': 'vae',
|
|
'name': metadata['vae'],
|
|
'hash': '',
|
|
'weight': None,
|
|
'required': False
|
|
}
|
|
dependencies.append(dep)
|
|
|
|
# ControlNet
|
|
if 'controlnet' in metadata and isinstance(metadata['controlnet'], list):
|
|
for controlnet in metadata['controlnet']:
|
|
if isinstance(controlnet, dict) and controlnet.get('model'):
|
|
dep = {
|
|
'type': 'controlnet',
|
|
'name': controlnet['model'],
|
|
'hash': '',
|
|
'weight': controlnet.get('weight', 1.0),
|
|
'required': False
|
|
}
|
|
dependencies.append(dep)
|
|
|
|
return dependencies
|
|
|
|
def _format_detailed_json(self, dependencies: list[dict], include_confidence: bool) -> str:
|
|
"""Format as detailed JSON"""
|
|
output = {
|
|
'metaman_dependencies': {
|
|
'version': '1.0.0',
|
|
'generated_at': datetime.now().isoformat(),
|
|
'total_dependencies': len(dependencies),
|
|
'found_count': len([d for d in dependencies if d.get('search_status') == 'found']),
|
|
},
|
|
'dependencies': dependencies
|
|
}
|
|
|
|
if not include_confidence:
|
|
# Remove confidence scores
|
|
for dep in output['dependencies']:
|
|
for source in dep.get('sources', []):
|
|
source.pop('confidence', None)
|
|
|
|
return json.dumps(output, indent=2, ensure_ascii=False)
|
|
|
|
def _format_download_urls(self, dependencies: list[dict]) -> str:
|
|
"""Format as simple download URL list"""
|
|
urls = []
|
|
|
|
for dep in dependencies:
|
|
for source in dep.get('sources', []):
|
|
if source.get('download_url'):
|
|
urls.append(f"{dep['name']}: {source['download_url']}")
|
|
|
|
return '\n'.join(urls) if urls else "No download URLs found"
|
|
|
|
def _format_summary_text(self, dependencies: list[dict]) -> str:
|
|
"""Format as human-readable summary"""
|
|
lines = []
|
|
lines.append("Model Dependencies Summary")
|
|
lines.append("=" * 30)
|
|
|
|
for dep in dependencies:
|
|
status = "✅" if dep.get('search_status') == 'found' else "❌"
|
|
lines.append(f"{status} {dep['type'].title()}: {dep['name']}")
|
|
|
|
if dep.get('hash'):
|
|
lines.append(f" Hash: {dep['hash'][:16]}...")
|
|
|
|
if dep.get('weight') is not None:
|
|
lines.append(f" Weight: {dep['weight']}")
|
|
|
|
sources = dep.get('sources', [])
|
|
if sources:
|
|
lines.append(f" Sources found: {len(sources)}")
|
|
for i, source in enumerate(sources[:2], 1):
|
|
lines.append(f" {i}. {source['platform']}: {source.get('confidence', 0):.2f} confidence")
|
|
else:
|
|
lines.append(" No sources found")
|
|
|
|
lines.append("")
|
|
|
|
return '\n'.join(lines)
|
|
|
|
def _format_dependency_list(self, dependencies: list[dict]) -> str:
|
|
"""Format as simple dependency list"""
|
|
lines = []
|
|
|
|
for dep in dependencies:
|
|
line = f"{dep['type']}: {dep['name']}"
|
|
if dep.get('weight') is not None:
|
|
line += f" (weight: {dep['weight']})"
|
|
if dep.get('hash'):
|
|
line += f" [hash: {dep['hash'][:8]}...]"
|
|
lines.append(line)
|
|
|
|
return '\n'.join(lines)
|
|
|
|
def _generate_summary(self, dependencies: list[dict]) -> str:
|
|
"""Generate a brief summary"""
|
|
total = len(dependencies)
|
|
found = len([d for d in dependencies if d.get('search_status') == 'found'])
|
|
|
|
summary = f"Dependencies: {found}/{total} found"
|
|
|
|
# Break down by type
|
|
by_type = {}
|
|
for dep in dependencies:
|
|
dep_type = dep['type']
|
|
if dep_type not in by_type:
|
|
by_type[dep_type] = {'total': 0, 'found': 0}
|
|
by_type[dep_type]['total'] += 1
|
|
if dep.get('search_status') == 'found':
|
|
by_type[dep_type]['found'] += 1
|
|
|
|
type_summaries = []
|
|
for dep_type, counts in by_type.items():
|
|
type_summaries.append(f"{dep_type}: {counts['found']}/{counts['total']}")
|
|
|
|
if type_summaries:
|
|
summary += f" ({', '.join(type_summaries)})"
|
|
|
|
return summary
|