148 lines
6.2 KiB
Python
148 lines
6.2 KiB
Python
import json
|
|
import random
|
|
import logging
|
|
from typing import Dict, Any, Optional, Tuple, List
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class WorkflowBuilder:
|
|
"""Helper to manipulate ComfyUI workflow JSONs."""
|
|
|
|
def __init__(self, workflow_json: Dict[str, Any]):
|
|
self.workflow = workflow_json
|
|
# Create lookups
|
|
self._nodes = self.workflow
|
|
if "nodes" in self.workflow and isinstance(self.workflow["nodes"], list):
|
|
# Handle "graph" format vs "api" format if needed.
|
|
# But usually for API we stick to the {node_id: node_data} format.
|
|
# If input is graph format, it might need conversion or distinct handling.
|
|
# Assuming API format for now as that's what's sent to /prompt.
|
|
pass
|
|
|
|
@classmethod
|
|
def from_json_string(cls, json_str: str) -> 'WorkflowBuilder':
|
|
"""Load from JSON int."""
|
|
return cls(json.loads(json_str))
|
|
|
|
def get_workflow(self) -> Dict[str, Any]:
|
|
"""Get the current workflow dict."""
|
|
return self.workflow
|
|
|
|
def set_prompt(self, positive: str, negative: Optional[str] = None) -> None:
|
|
"""
|
|
Attempt to set positive and negative prompts.
|
|
Heuristics:
|
|
- Look for CLIPTextEncode nodes.
|
|
- Often one is connected to KSampler 'positive' and one to 'negative'.
|
|
- Or look for custom titles like 'Positive Prompt', 'Negative Prompt'.
|
|
"""
|
|
# Simple heuristic: Find CLIPTextEncode nodes
|
|
# If we have title/coloring, we can use that.
|
|
# Otherwise, we might need graph traversal to see what connects to KSampler.
|
|
|
|
# For MVP, let's assume standard ComfyUI structure or look for specific titles first
|
|
|
|
positive_node_id = self._find_node_by_title("Positive Prompt")
|
|
negative_node_id = self._find_node_by_title("Negative Prompt")
|
|
|
|
# Fallback: Find KSampler and trace back
|
|
if not positive_node_id or not negative_node_id:
|
|
ksampler_id, ksampler = self._find_node_by_class("KSampler")
|
|
if ksampler:
|
|
# KSampler inputs: model, positive, negative, latent_image
|
|
if not positive_node_id:
|
|
positive_node_id = self._trace_input(ksampler, "positive")
|
|
if not negative_node_id:
|
|
negative_node_id = self._trace_input(ksampler, "negative")
|
|
|
|
if positive_node_id:
|
|
self._update_node_input(positive_node_id, "text", positive)
|
|
else:
|
|
logger.warning("Could not identify Positive Prompt node.")
|
|
|
|
if negative and negative_node_id:
|
|
self._update_node_input(negative_node_id, "text", negative)
|
|
elif negative:
|
|
logger.warning("Could not identify Negative Prompt node.")
|
|
|
|
def set_seed(self, seed: int) -> int:
|
|
"""Set seed on KSampler nodes or Seed nodes."""
|
|
# Find KSampler or anything with a 'seed' widget
|
|
updated = False
|
|
for node_id, node in self.workflow.items():
|
|
if "inputs" in node:
|
|
if "seed" in node["inputs"]:
|
|
# Ensure it's an int widget, not a link
|
|
if isinstance(node["inputs"]["seed"], (int, float)) or (isinstance(node["inputs"]["seed"], str) and node["inputs"]["seed"].isdigit()):
|
|
node["inputs"]["seed"] = seed
|
|
updated = True
|
|
if "noise_seed" in node["inputs"]:
|
|
# Some nodes call it noise_seed
|
|
if isinstance(node["inputs"]["noise_seed"], (int, float)):
|
|
node["inputs"]["noise_seed"] = seed
|
|
updated = True
|
|
|
|
if not updated:
|
|
logger.warning("Could not find any seed inputs to update.")
|
|
return seed
|
|
|
|
def set_image_dimensions(self, width: int, height: int) -> None:
|
|
"""Set width and height on EmptyLatentImage nodes."""
|
|
node_id, _ = self._find_node_by_class("EmptyLatentImage")
|
|
if node_id:
|
|
self._update_node_input(node_id, "width", width)
|
|
self._update_node_input(node_id, "height", height)
|
|
|
|
def set_steps(self, steps: int) -> None:
|
|
"""Set steps on KSampler."""
|
|
ksampler_ids = self._find_nodes_by_class("KSampler")
|
|
for nid in ksampler_ids:
|
|
self._update_node_input(nid, "steps", steps)
|
|
|
|
def set_cfg(self, cfg: float) -> None:
|
|
"""Set CFG scale on KSampler."""
|
|
ksampler_ids = self._find_nodes_by_class("KSampler")
|
|
for nid in ksampler_ids:
|
|
self._update_node_input(nid, "cfg", cfg)
|
|
|
|
def _find_node_by_title(self, title: str) -> Optional[str]:
|
|
"""Find node by its custom title (`_meta.title`)."""
|
|
for node_id, node in self.workflow.items():
|
|
if "_meta" in node and node["_meta"].get("title") == title:
|
|
return node_id
|
|
return None
|
|
|
|
def _find_node_by_class(self, class_type: str) -> Tuple[Optional[str], Optional[Dict]]:
|
|
"""Find first node of a specific class type."""
|
|
for node_id, node in self.workflow.items():
|
|
if node.get("class_type") == class_type:
|
|
return node_id, node
|
|
return None, None
|
|
|
|
def _find_nodes_by_class(self, class_type: str) -> List[str]:
|
|
"""Find all nodes of a specific class type."""
|
|
ids = []
|
|
for node_id, node in self.workflow.items():
|
|
if node.get("class_type") == class_type:
|
|
ids.append(node_id)
|
|
return ids
|
|
|
|
def _trace_input(self, node: Dict, input_name: str) -> Optional[str]:
|
|
"""
|
|
Trace back an input link to find the source node.
|
|
Input format in API JSON: "input_name": ["source_node_id", slot_index]
|
|
"""
|
|
if "inputs" not in node or input_name not in node["inputs"]:
|
|
return None
|
|
|
|
link = node["inputs"][input_name]
|
|
# Link structure: [node_id, slot_idx]
|
|
if isinstance(link, list) and len(link) == 2:
|
|
return str(link[0])
|
|
return None
|
|
|
|
def _update_node_input(self, node_id: str, input_name: str, value: Any) -> None:
|
|
if node_id in self.workflow and "inputs" in self.workflow[node_id]:
|
|
self.workflow[node_id]["inputs"][input_name] = value
|
|
|