mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-23 06:02:11 -03:00
362 lines
13 KiB
Python
362 lines
13 KiB
Python
"""
|
|
Node mappers for ComfyUI workflow parsing
|
|
"""
|
|
import logging
|
|
import os
|
|
import importlib.util
|
|
import inspect
|
|
from typing import Dict, List, Any, Optional, Union, Type, Callable, Tuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Global mapper registry
|
|
_MAPPER_REGISTRY: Dict[str, Dict] = {}
|
|
|
|
# =============================================================================
|
|
# Mapper Definition Functions
|
|
# =============================================================================
|
|
|
|
def create_mapper(
|
|
node_type: str,
|
|
inputs_to_track: List[str],
|
|
transform_func: Callable[[Dict], Any] = None
|
|
) -> Dict:
|
|
"""Create a mapper definition for a node type"""
|
|
mapper = {
|
|
"node_type": node_type,
|
|
"inputs_to_track": inputs_to_track,
|
|
"transform": transform_func or (lambda inputs: inputs)
|
|
}
|
|
return mapper
|
|
|
|
def register_mapper(mapper: Dict) -> None:
|
|
"""Register a node mapper in the global registry"""
|
|
_MAPPER_REGISTRY[mapper["node_type"]] = mapper
|
|
logger.debug(f"Registered mapper for node type: {mapper['node_type']}")
|
|
|
|
def get_mapper(node_type: str) -> Optional[Dict]:
|
|
"""Get a mapper for the specified node type"""
|
|
return _MAPPER_REGISTRY.get(node_type)
|
|
|
|
def get_all_mappers() -> Dict[str, Dict]:
|
|
"""Get all registered mappers"""
|
|
return _MAPPER_REGISTRY.copy()
|
|
|
|
# =============================================================================
|
|
# Node Processing Function
|
|
# =============================================================================
|
|
|
|
def process_node(node_id: str, node_data: Dict, workflow: Dict, parser: 'WorkflowParser') -> Any:
|
|
"""Process a node using its mapper and extract relevant information"""
|
|
node_type = node_data.get("class_type")
|
|
mapper = get_mapper(node_type)
|
|
|
|
if not mapper:
|
|
return None
|
|
|
|
result = {}
|
|
|
|
# Extract inputs based on the mapper's tracked inputs
|
|
for input_name in mapper["inputs_to_track"]:
|
|
if input_name in node_data.get("inputs", {}):
|
|
input_value = node_data["inputs"][input_name]
|
|
|
|
# Check if input is a reference to another node's output
|
|
if isinstance(input_value, list) and len(input_value) == 2:
|
|
try:
|
|
# Format is [node_id, output_slot]
|
|
ref_node_id, output_slot = input_value
|
|
# Convert node_id to string if it's an integer
|
|
if isinstance(ref_node_id, int):
|
|
ref_node_id = str(ref_node_id)
|
|
|
|
# Recursively process the referenced node
|
|
ref_value = parser.process_node(ref_node_id, workflow)
|
|
|
|
if ref_value is not None:
|
|
result[input_name] = ref_value
|
|
else:
|
|
# If we couldn't get a value from the reference, store the raw value
|
|
result[input_name] = input_value
|
|
except Exception as e:
|
|
logger.error(f"Error processing reference in node {node_id}, input {input_name}: {e}")
|
|
result[input_name] = input_value
|
|
else:
|
|
# Direct value
|
|
result[input_name] = input_value
|
|
|
|
# Apply the transform function
|
|
try:
|
|
return mapper["transform"](result)
|
|
except Exception as e:
|
|
logger.error(f"Error in transform function for node {node_id} of type {node_type}: {e}")
|
|
return result
|
|
|
|
# =============================================================================
|
|
# Default Mapper Definitions
|
|
# =============================================================================
|
|
|
|
def register_default_mappers() -> None:
|
|
"""Register all default mappers"""
|
|
|
|
# KSampler mapper
|
|
def transform_ksampler(inputs: Dict) -> Dict:
|
|
result = {
|
|
"seed": str(inputs.get("seed", "")),
|
|
"steps": str(inputs.get("steps", "")),
|
|
"cfg": str(inputs.get("cfg", "")),
|
|
"sampler": inputs.get("sampler_name", ""),
|
|
"scheduler": inputs.get("scheduler", ""),
|
|
}
|
|
|
|
# Process positive prompt
|
|
if "positive" in inputs:
|
|
result["prompt"] = inputs["positive"]
|
|
|
|
# Process negative prompt
|
|
if "negative" in inputs:
|
|
result["negative_prompt"] = inputs["negative"]
|
|
|
|
# Get dimensions from latent image
|
|
if "latent_image" in inputs and isinstance(inputs["latent_image"], dict):
|
|
width = inputs["latent_image"].get("width", 0)
|
|
height = inputs["latent_image"].get("height", 0)
|
|
if width and height:
|
|
result["size"] = f"{width}x{height}"
|
|
|
|
# Add clip_skip if present
|
|
if "clip_skip" in inputs:
|
|
result["clip_skip"] = str(inputs.get("clip_skip", ""))
|
|
|
|
return result
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="KSampler",
|
|
inputs_to_track=["seed", "steps", "cfg", "sampler_name", "scheduler",
|
|
"denoise", "positive", "negative", "latent_image",
|
|
"model", "clip_skip"],
|
|
transform_func=transform_ksampler
|
|
))
|
|
|
|
# EmptyLatentImage mapper
|
|
def transform_empty_latent(inputs: Dict) -> Dict:
|
|
width = inputs.get("width", 0)
|
|
height = inputs.get("height", 0)
|
|
return {"width": width, "height": height, "size": f"{width}x{height}"}
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="EmptyLatentImage",
|
|
inputs_to_track=["width", "height", "batch_size"],
|
|
transform_func=transform_empty_latent
|
|
))
|
|
|
|
# SD3LatentImage mapper - reuses same transform function as EmptyLatentImage
|
|
register_mapper(create_mapper(
|
|
node_type="EmptySD3LatentImage",
|
|
inputs_to_track=["width", "height", "batch_size"],
|
|
transform_func=transform_empty_latent
|
|
))
|
|
|
|
# CLIPTextEncode mapper
|
|
def transform_clip_text(inputs: Dict) -> Any:
|
|
return inputs.get("text", "")
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="CLIPTextEncode",
|
|
inputs_to_track=["text", "clip"],
|
|
transform_func=transform_clip_text
|
|
))
|
|
|
|
# LoraLoader mapper
|
|
def transform_lora_loader(inputs: Dict) -> Dict:
|
|
loras_data = inputs.get("loras", [])
|
|
lora_stack = inputs.get("lora_stack", {}).get("lora_stack", [])
|
|
|
|
lora_texts = []
|
|
|
|
# Process loras array
|
|
if isinstance(loras_data, dict) and "__value__" in loras_data:
|
|
loras_list = loras_data["__value__"]
|
|
elif isinstance(loras_data, list):
|
|
loras_list = loras_data
|
|
else:
|
|
loras_list = []
|
|
|
|
# Process each active lora entry
|
|
for lora in loras_list:
|
|
if isinstance(lora, dict) and lora.get("active", False):
|
|
lora_name = lora.get("name", "")
|
|
strength = lora.get("strength", 1.0)
|
|
lora_texts.append(f"<lora:{lora_name}:{strength}>")
|
|
|
|
# Process lora_stack if valid
|
|
if lora_stack and isinstance(lora_stack, list):
|
|
if not (len(lora_stack) == 2 and isinstance(lora_stack[0], (str, int)) and isinstance(lora_stack[1], int)):
|
|
for stack_entry in lora_stack:
|
|
lora_name = stack_entry[0]
|
|
strength = stack_entry[1]
|
|
lora_texts.append(f"<lora:{lora_name}:{strength}>")
|
|
|
|
return {"loras": " ".join(lora_texts)}
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="Lora Loader (LoraManager)",
|
|
inputs_to_track=["loras", "lora_stack"],
|
|
transform_func=transform_lora_loader
|
|
))
|
|
|
|
# LoraStacker mapper
|
|
def transform_lora_stacker(inputs: Dict) -> Dict:
|
|
loras_data = inputs.get("loras", [])
|
|
result_stack = []
|
|
|
|
# Handle existing stack entries
|
|
existing_stack = []
|
|
lora_stack_input = inputs.get("lora_stack", [])
|
|
|
|
if isinstance(lora_stack_input, dict) and "lora_stack" in lora_stack_input:
|
|
existing_stack = lora_stack_input["lora_stack"]
|
|
elif isinstance(lora_stack_input, list):
|
|
if not (len(lora_stack_input) == 2 and isinstance(lora_stack_input[0], (str, int)) and
|
|
isinstance(lora_stack_input[1], int)):
|
|
existing_stack = lora_stack_input
|
|
|
|
# Add existing entries
|
|
if existing_stack:
|
|
result_stack.extend(existing_stack)
|
|
|
|
# Process new loras
|
|
if isinstance(loras_data, dict) and "__value__" in loras_data:
|
|
loras_list = loras_data["__value__"]
|
|
elif isinstance(loras_data, list):
|
|
loras_list = loras_data
|
|
else:
|
|
loras_list = []
|
|
|
|
for lora in loras_list:
|
|
if isinstance(lora, dict) and lora.get("active", False):
|
|
lora_name = lora.get("name", "")
|
|
strength = float(lora.get("strength", 1.0))
|
|
result_stack.append((lora_name, strength))
|
|
|
|
return {"lora_stack": result_stack}
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="Lora Stacker (LoraManager)",
|
|
inputs_to_track=["loras", "lora_stack"],
|
|
transform_func=transform_lora_stacker
|
|
))
|
|
|
|
# JoinStrings mapper
|
|
def transform_join_strings(inputs: Dict) -> str:
|
|
string1 = inputs.get("string1", "")
|
|
string2 = inputs.get("string2", "")
|
|
delimiter = inputs.get("delimiter", "")
|
|
return f"{string1}{delimiter}{string2}"
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="JoinStrings",
|
|
inputs_to_track=["string1", "string2", "delimiter"],
|
|
transform_func=transform_join_strings
|
|
))
|
|
|
|
# StringConstant mapper
|
|
def transform_string_constant(inputs: Dict) -> str:
|
|
return inputs.get("string", "")
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="StringConstantMultiline",
|
|
inputs_to_track=["string"],
|
|
transform_func=transform_string_constant
|
|
))
|
|
|
|
# TriggerWordToggle mapper
|
|
def transform_trigger_word_toggle(inputs: Dict) -> str:
|
|
toggle_data = inputs.get("toggle_trigger_words", [])
|
|
|
|
if isinstance(toggle_data, dict) and "__value__" in toggle_data:
|
|
toggle_words = toggle_data["__value__"]
|
|
elif isinstance(toggle_data, list):
|
|
toggle_words = toggle_data
|
|
else:
|
|
toggle_words = []
|
|
|
|
# Filter active trigger words
|
|
active_words = []
|
|
for item in toggle_words:
|
|
if isinstance(item, dict) and item.get("active", False):
|
|
word = item.get("text", "")
|
|
if word and not word.startswith("__dummy"):
|
|
active_words.append(word)
|
|
|
|
return ", ".join(active_words)
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="TriggerWord Toggle (LoraManager)",
|
|
inputs_to_track=["toggle_trigger_words"],
|
|
transform_func=transform_trigger_word_toggle
|
|
))
|
|
|
|
# FluxGuidance mapper
|
|
def transform_flux_guidance(inputs: Dict) -> Dict:
|
|
result = {}
|
|
|
|
if "guidance" in inputs:
|
|
result["guidance"] = inputs["guidance"]
|
|
|
|
if "conditioning" in inputs:
|
|
conditioning = inputs["conditioning"]
|
|
if isinstance(conditioning, str):
|
|
result["prompt"] = conditioning
|
|
else:
|
|
result["prompt"] = "Unknown prompt"
|
|
|
|
return result
|
|
|
|
register_mapper(create_mapper(
|
|
node_type="FluxGuidance",
|
|
inputs_to_track=["guidance", "conditioning"],
|
|
transform_func=transform_flux_guidance
|
|
))
|
|
|
|
# =============================================================================
|
|
# Extension Loading
|
|
# =============================================================================
|
|
|
|
def load_extensions(ext_dir: str = None) -> None:
|
|
"""
|
|
Load mapper extensions from the specified directory
|
|
|
|
Extension files should define mappers using the create_mapper function
|
|
and then call register_mapper to add them to the registry.
|
|
"""
|
|
# Use default path if none provided
|
|
if ext_dir is None:
|
|
# Get the directory of this file
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
ext_dir = os.path.join(current_dir, 'ext')
|
|
|
|
# Ensure the extension directory exists
|
|
if not os.path.exists(ext_dir):
|
|
os.makedirs(ext_dir, exist_ok=True)
|
|
logger.info(f"Created extension directory: {ext_dir}")
|
|
return
|
|
|
|
# Load each Python file in the extension directory
|
|
for filename in os.listdir(ext_dir):
|
|
if filename.endswith('.py') and not filename.startswith('_'):
|
|
module_path = os.path.join(ext_dir, filename)
|
|
module_name = f"workflow.ext.{filename[:-3]}" # Remove .py
|
|
|
|
try:
|
|
# Load the module
|
|
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
if spec and spec.loader:
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
logger.info(f"Loaded extension module: {filename}")
|
|
except Exception as e:
|
|
logger.warning(f"Error loading extension {filename}: {e}")
|
|
|
|
# Initialize the registry with default mappers
|
|
register_default_mappers() |