Martin Rizzo
d1170f9d96
This commit refactors the `bulkgen.py` script to improve code organization and readability. Comments have been added to key classes and methods to clarify logic and reduce complexity, making the codebase easier to follow and maintain.
764 lines
30 KiB
Python
764 lines
30 KiB
Python
"""
|
|
File : bulkgen.py
|
|
Purpose : Script to create a batch of images based on a configuration file.
|
|
Author : Martin Rizzo | <martinrizzo@gmail.com>
|
|
Date : Mar 14, 2026
|
|
Repo : https://github.com/martin-rizzo/ComfyUI-ZImagePowerNodes
|
|
License : MIT
|
|
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
|
|
ComfyUI-ZImagePowerNodes
|
|
ComfyUI nodes designed specifically for the "Z-Image" model.
|
|
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
|
|
"""
|
|
import os
|
|
import sys
|
|
import argparse
|
|
from typing import overload
|
|
from copy import deepcopy
|
|
|
|
import json
|
|
import urllib.request
|
|
import urllib.parse
|
|
import uuid
|
|
|
|
# requires: pip install websocket-client
|
|
# or: sudo dnf install python3-websocket-client
|
|
import websocket
|
|
|
|
|
|
DEFAULT_HOST = os.getenv("BULKGEN_COMFY_HOST") or os.getenv("COMFY_HOST") or "127.0.0.1"
|
|
DEFAULT_PORT = os.getenv("BULKGEN_COMFY_PORT") or os.getenv("COMFY_PORT") or "8188"
|
|
DEFAULT_PROMPT =""
|
|
DEFAULT_CHECKPOINT =os.getenv("BULKGEN_CHECKPOINT" ) or "z_image_turbo-Q5_K_S.gguf"
|
|
DEFAULT_TEXT_ENCODER =os.getenv("BULKGEN_TEXT_ENCODER" ) or "Qwen3-4B.i1-Q5_K_S.gguf"
|
|
DEFAULT_STYLE =os.getenv("BULKGEN_STYLE" ) or "Casual Photo"
|
|
DEFAULT_STEPS =os.getenv("BULKGEN_STEPS" ) or "9"
|
|
DEFAULT_SEED =os.getenv("BULKGEN_SEED" ) or "1"
|
|
DEFAULT_INCALIBRATION=os.getenv("BULKGEN_INCALIBRATION") or "1.0"
|
|
DEFAULT_LOWRES_BIAS = False
|
|
DEFAULT_DENOISE = 1.0
|
|
README_NOTE ="This workflow was generated by bulkgen.py"
|
|
|
|
FILE_PREFIX = "bgen"
|
|
|
|
|
|
# get the directory where the script is located
|
|
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
# ANSI escape codes for colored terminal output
|
|
RED = '\033[91m'
|
|
DKRED = '\033[31m'
|
|
YELLOW = '\033[93m'
|
|
DKYELLOW = '\033[33m'
|
|
GREEN = '\033[92m'
|
|
CYAN = '\033[96m'
|
|
DKGRAY = '\033[90m'
|
|
RESET = '\033[0m'
|
|
|
|
#============================= ERROR MESSAGES ==============================#
|
|
|
|
def disable_colors():
|
|
global RED, DKRED, YELLOW, DKYELLOW, GREEN, CYAN, DKGRAY, RESET
|
|
RED, DKRED, YELLOW, DKYELLOW, GREEN, CYAN, DKGRAY, RESET = "", "", "", "", "", "", "", ""
|
|
|
|
|
|
def info(message: str, padding: int = 0, file=sys.stderr) -> None:
|
|
"""Displays an informational message to the error stream.
|
|
"""
|
|
print(f"{" "*padding}{CYAN}\u24d8 {message}{RESET}", file=file)
|
|
|
|
|
|
def warning(message: str, *info_messages: str, padding: int = 0, file=sys.stderr) -> None:
|
|
"""Displays a warning message to the standard error stream.
|
|
"""
|
|
print(f"{" "*padding}{CYAN}[{YELLOW}WARNING{CYAN}]{DKYELLOW} {message}{RESET}", file=file)
|
|
for info_message in info_messages:
|
|
info(info_message, padding=padding, file=file)
|
|
|
|
|
|
def error(message: str, *info_messages: str, padding: int = 0, file=sys.stderr) -> None:
|
|
"""Displays an error message to the standard error stream.
|
|
"""
|
|
print(f"{" "*padding}{DKRED}[{RED}ERROR!{DKRED}]{DKYELLOW} {message}{RESET}", file=file)
|
|
for info_message in info_messages:
|
|
info(info_message, padding=padding, file=file)
|
|
|
|
|
|
def fatal_error(message: str, *info_messages: str, padding: int = 0, file=sys.stderr) -> None:
|
|
"""Displays a fatal error message to the standard error stream and exits with status code 1.
|
|
"""
|
|
error(message, *info_messages, padding=padding, file=file)
|
|
sys.exit(1)
|
|
|
|
|
|
#================================= HELPERS =================================#
|
|
|
|
def _to_float(value: str) -> float:
|
|
"""Converts a string to float, handling optional percentage signs."""
|
|
value = value.strip()
|
|
if not value: raise ValueError("Empty string provided for float conversion.")
|
|
elif value.endswith('%') : return float(value[:-1]) / 100.0
|
|
elif value.startswith('%'): return float(value[1:] ) / 100.0
|
|
return float(value)
|
|
|
|
|
|
def _to_int(value: str) -> int:
|
|
"""Converts a string to int."""
|
|
value = value.strip()
|
|
if not value: raise ValueError("Empty string provided for int conversion.")
|
|
return int(float(value))
|
|
|
|
|
|
@overload
|
|
def to_float(values: str) -> float: ...
|
|
|
|
@overload
|
|
def to_float(values: list[str]) -> list[float]: ...
|
|
|
|
def to_float(values: str | list[str]) -> float | list[float]:
|
|
"""Converts a string (or list of strings) to float(s), supporting percentages.
|
|
Raises:
|
|
ValueError: If the input contains invalid characters for numbers/percentages.
|
|
"""
|
|
try:
|
|
if isinstance(values, str):
|
|
return _to_float(values)
|
|
return [_to_float(value) for value in values]
|
|
except ValueError as e:
|
|
raise ValueError(f"Invalid float/percent values: {values}") from e
|
|
|
|
|
|
@overload
|
|
def to_int(values: str) -> int: ...
|
|
|
|
@overload
|
|
def to_int(values: list[str]) -> list[int]: ...
|
|
|
|
def to_int(values: str | list[str]) -> int | list[int]:
|
|
"""Converts a string (or list of strings) to int(s).
|
|
Raises:
|
|
ValueError: If the input contains invalid characters for integer numbers.
|
|
"""
|
|
try:
|
|
if isinstance(values, str):
|
|
return _to_int(values)
|
|
return [_to_int(value) for value in values]
|
|
except (ValueError, TypeError) as e:
|
|
raise ValueError(f"Invalid int values: {values}") from e
|
|
|
|
|
|
def node_has_input(node: dict, input_name: str) -> bool:
|
|
"""Return `True` if the node has an input with the given name, (case insensitive)
|
|
"""
|
|
lower_input_name = input_name.lower()
|
|
for input in node.get("inputs", []):
|
|
if input.get("name","").lower() == lower_input_name:
|
|
return True
|
|
return False
|
|
|
|
|
|
#============================== FIXING VALUES ==============================#
|
|
|
|
def fix_percent(value: float | int, use_off: bool = False) -> str:
|
|
if use_off and value<=0:
|
|
return "off"
|
|
percent_str = f"{int(value*100)}%"
|
|
return percent_str
|
|
|
|
|
|
#================================ GenParams ================================#
|
|
|
|
class GenParams:
|
|
"""
|
|
Represents a generation parameter configuration (with multiple optional values)
|
|
|
|
This class allows storing multiple values for each parameter to support
|
|
combinatorial generation. For example, if you want to test different
|
|
checkpoints, distinct steps, and different seeds but with the same prompt.
|
|
|
|
Example Usage:
|
|
multi_params = GenParams(
|
|
prompt="A cyberpunk city",
|
|
style=["anime", "phone photo"]
|
|
checkpoint=["z_image_turbo-Q5_K_S.gguf", "z_image_turbo_bf16.safetensors"],
|
|
steps=[5,7,9],
|
|
seed=[1]
|
|
)
|
|
# Later logic could generate all images using `self.generate_combinations()`
|
|
|
|
Attributes:
|
|
prompt (str) : The single prompt text shared across all variations.
|
|
checkpoint (List[str]) : List of model checkpoint names.
|
|
text_encoder (List[str]) : List of text encoder names.
|
|
style (List[str]) : List of artistic styles to test.
|
|
steps (List[int]) : List of denoising step counts.
|
|
seed (List[int]) : List of random seeds.
|
|
incalibration (List[float]): List of initial noise calibration values.
|
|
filename (str) : Optional filename for the image generated by this parameter set.
|
|
"""
|
|
def __init__(self, *,
|
|
prompt : str | None = None,
|
|
checkpoint : list[str] | None = None,
|
|
text_encoder : list[str] | None = None,
|
|
style : list[str] | None = None,
|
|
steps : list[int] | None = None,
|
|
seed : list[int] | None = None,
|
|
incalibration: list[float] | None = None,
|
|
filename : str | None = None,
|
|
) -> None:
|
|
self.prompt = prompt if type(prompt) == str else DEFAULT_PROMPT
|
|
self.checkpoint = checkpoint if type(checkpoint) == list else [DEFAULT_CHECKPOINT]
|
|
self.text_encoder = text_encoder if type(text_encoder) == list else [DEFAULT_TEXT_ENCODER]
|
|
self.style = style if type(style) == list else [DEFAULT_STYLE]
|
|
self.steps = steps if type(steps) == list else [int(DEFAULT_STEPS)]
|
|
self.seed = seed if type(seed) == list else [int(DEFAULT_SEED)]
|
|
self.incalibration = incalibration if type(incalibration) == list else [float(DEFAULT_INCALIBRATION)]
|
|
self.filename = filename if type(filename) == str else ""
|
|
|
|
|
|
def set_param(self, param_and_values: str):
|
|
"""
|
|
Updates a generation parameter based on a string with format "param = value"
|
|
|
|
This method allows modifying parameters like checkpoint, style, steps, etc.,
|
|
using a key-value string format. Param names can be prefixed with "$".
|
|
|
|
Args:
|
|
param_and_values (str): The parameter name and values in the format "param=value"
|
|
or "param=value1,value2,...".
|
|
Example: "ckpt=zimage.safetensors" or "$steps=5,7,9".
|
|
Raises:
|
|
ValueError: If the parameter name is not recognized (e.g., unknown param).
|
|
ValueError: If value conversion fails (handled by to_int/to_float).
|
|
|
|
Supported Parameters:
|
|
- "ckpt" : Updates self.checkpoint.
|
|
- "style": Updates self.style.
|
|
- "steps": Updates self.steps.
|
|
- "seed" : Updates self.seed .
|
|
- "inc", "incalibration", "initial_noise_calibration": Updates self.incalibration.
|
|
"""
|
|
param, _, values = param_and_values.partition("=")
|
|
param = param.strip().lower()
|
|
if param.startswith('$'): param = param[1:]
|
|
|
|
# builds the list of values
|
|
values = [v.strip() for v in values.split(",")]
|
|
|
|
if param == "ckpt":
|
|
self.checkpoint = values
|
|
elif param == "style":
|
|
self.style = values
|
|
elif param == "steps":
|
|
self.steps = to_int(values)
|
|
elif param == "seed":
|
|
self.seed = to_int(values)
|
|
elif param == "inc" or param == "incalibration" or param == "initial_noise_calibration":
|
|
self.incalibration = to_float(values)
|
|
else:
|
|
raise ValueError(f"Unknown generation parameter: {param}")
|
|
|
|
|
|
def generate_combinations(self, name_prefix="bulk") -> list["GenParams"]:
|
|
"""
|
|
Generates a list of GenParams objects representing every possible
|
|
combination of the multi-valued parameters defined in this instance.
|
|
|
|
Returns:
|
|
List[GenParams]: A list containing unique GenParams configurations.
|
|
|
|
Example:
|
|
params_list = self.generate_combinations()
|
|
"""
|
|
from itertools import product
|
|
result_list: list[GenParams] = []
|
|
|
|
# define the options to iterate over
|
|
options = [ self.checkpoint,
|
|
self.style,
|
|
self.steps,
|
|
self.seed,
|
|
self.incalibration ]
|
|
|
|
# create a new GenParams instance for each combination
|
|
for combo in list(product(*options)):
|
|
filename = name_prefix
|
|
if len(self.steps)>1: filename += f"_steps{combo[2]}"
|
|
if len(self.seed )>1: filename += f"_seed{combo[3]}"
|
|
genparams = GenParams(
|
|
prompt = self.prompt,
|
|
filename = filename,
|
|
checkpoint = [ combo[0] ],
|
|
style = [ combo[1] ],
|
|
steps = [ combo[2] ],
|
|
seed = [ combo[3] ],
|
|
incalibration = [ combo[4] ],
|
|
)
|
|
result_list.append(genparams)
|
|
|
|
return result_list
|
|
|
|
|
|
def copy(self) -> 'GenParams':
|
|
"""Returns a deep copy with all attributes duplicated independently.
|
|
Example:
|
|
original = GenParams(prompt="A cyberpunk city", steps=[5, 7, 9])
|
|
copy_instance = original.copy()
|
|
# Modifying copy_instance.steps will not affect original
|
|
original.steps.append(10)
|
|
assert copy_instance.steps == [5, 7, 9] # True
|
|
"""
|
|
return deepcopy(self)
|
|
|
|
|
|
def __str__(self) -> str:
|
|
"""Returns a readable string representation of the multi-value configuration."""
|
|
return f"GenParams(checkpoint={self.checkpoint}, style={self.style}, " \
|
|
f"steps={self.steps}, seed={self.seed}, incalibration={self.incalibration}, "\
|
|
f"prompt=\"{self.prompt}\")"
|
|
|
|
|
|
#================================= CONFIG ==================================#
|
|
|
|
class Config:
|
|
"""
|
|
Class responsible for parsing and managing configuration files.
|
|
"""
|
|
def __init__(self):
|
|
self._action = ""
|
|
self._content = ""
|
|
self._is_first_line = True
|
|
self._is_shebang_line = False
|
|
self._is_comment_zone = True
|
|
self._global_params = GenParams()
|
|
self._custom_styles : dict[str, str] = {}
|
|
self._genparams_list : list[GenParams] = []
|
|
|
|
@classmethod
|
|
def from_file(cls, path: str) -> "Config":
|
|
"""Loads configuration from a file.
|
|
Args:
|
|
path (str): The file path to the configuration file.
|
|
Returns:
|
|
Config: The instantiated configuration object after loading.
|
|
"""
|
|
config = cls()
|
|
config.load_begin()
|
|
with open(path, "r") as file:
|
|
for line in file: config.load_text_line(line)
|
|
config.load_end()
|
|
return config
|
|
|
|
@property
|
|
def all_genparams(self) -> list[GenParams]:
|
|
"""Returns the list of all generation parameters stored in the configuration."""
|
|
return self._genparams_list
|
|
|
|
@property
|
|
def custom_styles(self) -> dict[str, str]:
|
|
"""Returns the dictionary of custom styles defined in the configuration."""
|
|
return self._custom_styles
|
|
|
|
|
|
def load_begin(self):
|
|
"""Initializes the loading process."""
|
|
pass
|
|
|
|
|
|
def load_text_line(self, line: str):
|
|
"""Loads a text line from the configuration file
|
|
|
|
The configuration loading process should be done line by line, wrapping
|
|
the whole process with load_begin() and load_end() methods.
|
|
"""
|
|
self._is_shebang_line = self._is_first_line and line.startswith("#!")
|
|
self._is_first_line = False
|
|
|
|
# note: trailing whitespaces are lost at the end of each line
|
|
line = line.rstrip()
|
|
|
|
# discard comments if it is in a zone that can have comments
|
|
if self._is_comment_zone and line.startswith(">>#"):
|
|
return
|
|
|
|
# verify if the line is a new action to be processed
|
|
if ( self._is_shebang_line or #< sheban "#!ZCONFIG" (compatibility with Amazing Z-Image Workflow)
|
|
line.startswith(">::") or #< action to modify workflow (compatibility with Amazing Z-Image Workflow)
|
|
line.startswith(">>>") or #< style selection / prompt start
|
|
line.startswith(">>$") #< parameter configuration
|
|
):
|
|
self._process_config_action(self._action, self._content)
|
|
self._action, self._content = line, ""
|
|
else:
|
|
self._content += line + "\n"
|
|
|
|
|
|
def load_end(self):
|
|
"""Finalizes the loading process."""
|
|
self._process_config_action(self._action, self._content)
|
|
self._action, self._content = "", ""
|
|
|
|
|
|
def add_new_genparams(self, genparams: GenParams):
|
|
self._genparams_list.append( genparams )
|
|
|
|
|
|
def _process_config_action(self, action: str, content: str, /):
|
|
"""
|
|
Process configuration commands based on the action syntax.
|
|
|
|
Parses the provided action to determine if it is a parameter assignment, a
|
|
prompt definition with parameters, or a simple prompt block, and executes
|
|
the corresponding logic.
|
|
|
|
Args:
|
|
action (str): The action to be processed (case inseensitive)
|
|
content (str): The accompanying content associated with the action.
|
|
"""
|
|
action = action.strip()
|
|
action_lower = action.lower()
|
|
content = content.strip()
|
|
|
|
# the action of ">>>Prompt" can have an optional ':' on the right side
|
|
if action_lower.startswith(">>>prompt"):
|
|
action = action.rstrip(":")
|
|
action_lower = action_lower.rstrip(":")
|
|
|
|
# processes global parameter assignment
|
|
# format: ">>>$param = valueA,valueB,...,valueN"
|
|
if action.startswith(">>>$"):
|
|
self._global_params.set_param( action[3:] )
|
|
|
|
# processes a prompt with parameters
|
|
# format: ">>>prompt{ param=valueA,valueB,...; param2=value2; ... }"
|
|
# and following lines are the full prompt
|
|
elif action_lower.startswith(">>>prompt{") and action.endswith("}"):
|
|
genparams = self._global_params.copy()
|
|
genparams.prompt = content.strip()
|
|
for param_and_values in action[10:-1].split(";"):
|
|
genparams.set_param( param_and_values )
|
|
self.add_new_genparams(genparams)
|
|
|
|
# processes prompt without parameters
|
|
# format: ">>>" or ">>>prompt:"
|
|
# and following lines are the full prompt
|
|
elif action_lower == ">>>prompt" or action == ">>>":
|
|
genparams = self._global_params.copy()
|
|
genparams.prompt = content.strip()
|
|
self.add_new_genparams(genparams)
|
|
|
|
# processes classical style selection
|
|
# format: ">>>Style Name"
|
|
elif action.startswith(">>>"):
|
|
style = action[3:].strip()
|
|
self._global_params.set_param(f"style={style}")
|
|
if content:
|
|
# TODO: verify that there are no commas
|
|
self._custom_styles[style] = content
|
|
|
|
|
|
|
|
|
|
#================================ ComfyAPI =================================#
|
|
|
|
class ComfyAPI:
|
|
|
|
def __init__(self, host="127.0.0.1", port="8188"):
|
|
self.server_address = f"{host}:{port}"
|
|
self.client_id = str(uuid.uuid4())
|
|
|
|
def _send_request(self, endpoint, data=None):
|
|
url = f"http://{self.server_address}{endpoint}"
|
|
headers = {'Content-Type': 'application/json'}
|
|
|
|
payload = json.dumps(data).encode('utf-8') if data else None
|
|
req = urllib.request.Request(url, data=payload, headers=headers)
|
|
|
|
with urllib.request.urlopen(req) as response:
|
|
return json.loads(response.read().decode('utf-8'))
|
|
|
|
|
|
def check_connection(self):
|
|
"""Returns True if a connection to the server can be established."""
|
|
try:
|
|
url = f"http://{self.server_address}/"
|
|
req = urllib.request.Request(url)
|
|
with urllib.request.urlopen(req, timeout=5) as response:
|
|
if 200 <= response.status < 300:
|
|
return True
|
|
else:
|
|
return False
|
|
except:
|
|
return False
|
|
|
|
|
|
def get_history(self, prompt_id):
|
|
return self._send_request(f"/history/{prompt_id}")
|
|
|
|
def download_image(self, filename, subfolder, folder_type):
|
|
params = urllib.parse.urlencode({
|
|
"filename": filename,
|
|
"subfolder": subfolder,
|
|
"type": folder_type
|
|
})
|
|
url = f"http://{self.server_address}/view?{params}"
|
|
with urllib.request.urlopen(url) as response:
|
|
return response.read()
|
|
|
|
def generate_and_wait(self, workflow, show_progress=True):
|
|
ws = websocket.WebSocket()
|
|
ws.connect(f"ws://{self.server_address}/ws?clientId={self.client_id}")
|
|
|
|
# Queue the prompt
|
|
queue_data = {"prompt": workflow, "client_id": self.client_id}
|
|
prompt_id = self._send_request("/prompt", queue_data)['prompt_id']
|
|
|
|
while True:
|
|
out = ws.recv()
|
|
if isinstance(out, str):
|
|
msg = json.loads(out)
|
|
|
|
# Update Progress
|
|
if show_progress and msg['type'] == 'progress':
|
|
val = msg['data']['value']
|
|
max_val = msg['data']['max']
|
|
percent = int(100 * val / max_val)
|
|
bar = "█" * (percent // 2) + "-" * (50 - percent // 2)
|
|
print(f"\rGenerating: |{bar}| {percent}%", end="", flush=True)
|
|
|
|
# Check if finished
|
|
if msg['type'] == 'executing':
|
|
if msg['data']['node'] is None and msg['data']['prompt_id'] == prompt_id:
|
|
print()
|
|
break
|
|
|
|
# Fetch output info
|
|
history = self.get_history(prompt_id)[prompt_id]
|
|
images_binary = []
|
|
|
|
for node_id in history['outputs']:
|
|
node_output = history['outputs'][node_id]
|
|
if 'images' in node_output:
|
|
for img in node_output['images']:
|
|
img_data = self.download_image(img['filename'], img['subfolder'], img['type'])
|
|
images_binary.append(img_data)
|
|
|
|
return images_binary
|
|
|
|
|
|
#============================== IMAGE BUILDER ==============================#
|
|
|
|
class ImageBuilder:
|
|
"""
|
|
Generates images through the ComfyUI API.
|
|
"""
|
|
def __init__(self, comfyapi: ComfyAPI) -> None:
|
|
self._confyapi = comfyapi
|
|
self._last_checkpoint = ""
|
|
|
|
|
|
def __call__(self, genparams: GenParams) -> bytes:
|
|
"""Convenience alias for `build_images()`"""
|
|
return self.build_image(genparams)
|
|
|
|
|
|
def build_image(self, genparams: GenParams) -> bytes:
|
|
"""
|
|
Generates images according to the provided generation parameters.
|
|
|
|
Args:
|
|
genparams: Object containing generation parameters (style, prompt, seed, etc.).
|
|
comfyapi : The API client instance for interacting with ComfyUI.
|
|
Returns:
|
|
List of generated images (raw bytes).
|
|
"""
|
|
# notify user when the checkpoint changes
|
|
if self._last_checkpoint != genparams.checkpoint[0]:
|
|
self._last_checkpoint = genparams.checkpoint[0]
|
|
info(f"Loading model \"{genparams.checkpoint[0]}\", this may take a while...")
|
|
|
|
workflow = self.build_workflow(genparams)
|
|
api_workflow = self.build_api_workflow(genparams)
|
|
return self._confyapi.generate_and_wait(api_workflow)[0]
|
|
|
|
|
|
def build_workflow(self, genparams: GenParams, custom_styles: str = "") -> dict:
|
|
"""
|
|
Builds a ComfyUI workflow by loading a template JSON and injecting parameters.
|
|
|
|
Loads the "bulkgen.json" workflow file, iterates through nodes, and
|
|
updates widget values based oon the provided generation parameters.
|
|
|
|
Args:
|
|
genparams : A `GenParams` object containing generation parameters (style, prompt, seed, etc.).
|
|
custom_styles: Optional custom style templates.
|
|
Returns:
|
|
The modified workflow dictionary with updated generation parameters.
|
|
"""
|
|
workflow_file = SCRIPT_DIR + "/bulkgen.json"
|
|
with open(workflow_file, 'r') as f:
|
|
workflow = json.load(f)
|
|
|
|
# get the workflow's list of nodes
|
|
nodes = workflow.get("nodes")
|
|
if not nodes: return {}
|
|
|
|
for node in nodes:
|
|
ntype : str = node.get("type","")
|
|
title : str = node.get("title","")
|
|
title_lower = title.lower()
|
|
|
|
# README
|
|
if ntype == "Note" and title_lower == "readme":
|
|
node["widgets_values"] = [ README_NOTE ]
|
|
|
|
# CUSTOM STYLES
|
|
if ntype == "PrimitiveStringMultiline" and "styles" in title_lower:
|
|
node["widgets_values"] = [ custom_styles ]
|
|
|
|
# Style + Prompt
|
|
if "StylePrompt" in ntype and node_has_input(node,"clip"):
|
|
number_of_widgets = len(node.get("widgets_values",[]))
|
|
if number_of_widgets != 4:
|
|
warning(f"Node \"{title}\" has {number_of_widgets} widgets, expected 4")
|
|
node["widgets_values"] = [
|
|
f'"{genparams.style[0]}"', None, None, genparams.prompt ]
|
|
|
|
# Seed + Steps + InCalibration
|
|
if "ZSamplerTurbo" in ntype and node_has_input(node,"model"):
|
|
number_of_widgets = len(node.get("widgets_values",[]))
|
|
if number_of_widgets != 7:
|
|
warning(f"Node \"{title}\" has {number_of_widgets} widgets, expected 7")
|
|
node["widgets_values"] = [
|
|
genparams.seed[0], "fixed", genparams.steps[0], DEFAULT_DENOISE, None,
|
|
fix_percent(genparams.incalibration[0], use_off=True),
|
|
DEFAULT_LOWRES_BIAS ]
|
|
|
|
# Checkpoint
|
|
if ntype == "UnetLoaderGGUF":
|
|
node["widgets_values"] = [
|
|
genparams.checkpoint[0] ]
|
|
|
|
# TextEncoder
|
|
if ntype == "CLIPLoaderGGUF":
|
|
node["widgets_values"] = [
|
|
genparams.text_encoder[0], "lumina2" ]
|
|
|
|
return workflow
|
|
|
|
|
|
def build_api_workflow(self, genparams: GenParams, custom_styles: str = "") -> dict:
|
|
"""
|
|
Builds a ComfyUI API workflow by loading a template JSON and updating inputs.
|
|
|
|
Loads the "bulkgen.api.json" workflow file, iterates through nodes, and
|
|
updates the input values based on the provided generation parameters.
|
|
|
|
Args:
|
|
genparams : A `GenParams` object containing generation parameters (style, prompt, seed, etc.).
|
|
custom_styles: Optional custom style templates.
|
|
Returns:
|
|
The modified API workflow dictionary, ready to be sent to "/prompt"
|
|
via the ComfyUI REST API.
|
|
"""
|
|
api_workflow_file = SCRIPT_DIR + "/bulkgen.api.json"
|
|
with open(api_workflow_file, 'r') as f:
|
|
api_workflow = json.load(f)
|
|
|
|
for _, node in api_workflow.items():
|
|
ntype : str = node.get("class_type","")
|
|
inputs : dict = node.get("inputs",{})
|
|
|
|
# Style + Prompt
|
|
if ntype.startswith("StylePrompt") and ntype.endswith("ZImagePowerNodes"):
|
|
inputs["style"] = genparams.style[0]
|
|
inputs["text" ] = genparams.prompt
|
|
|
|
# # Seed + Steps + InCalibration
|
|
if ntype.startswith("ZSamplerTurbo") and ntype.endswith("ZImagePowerNodes"):
|
|
inputs["seed" ] = genparams.seed[0]
|
|
inputs["steps" ] = genparams.steps[0]
|
|
inputs["initial_noise_calibration"] = fix_percent(genparams.incalibration[0], use_off=True)
|
|
inputs["denoise" ] = DEFAULT_DENOISE
|
|
inputs["lowres_bias"] = DEFAULT_LOWRES_BIAS
|
|
|
|
# Checkpoint
|
|
if ntype == "UnetLoaderGGUF":
|
|
inputs["unet_name"] = genparams.checkpoint[0]
|
|
|
|
# TextEncoder
|
|
if ntype == "CLIPLoaderGGUF":
|
|
inputs["clip_name"] = genparams.text_encoder[0]
|
|
|
|
return api_workflow
|
|
|
|
|
|
|
|
#===========================================================================#
|
|
#////////////////////////////////// MAIN ///////////////////////////////////#
|
|
#===========================================================================#
|
|
|
|
|
|
def main(args=None, parent_script=None):
|
|
"""
|
|
Main entry point for the script.
|
|
Args:
|
|
args (optional): List of arguments to parse. Default is None, which will use the command line arguments.
|
|
parent_script (optional): The name of the calling script if any. Used for customizing help output.
|
|
"""
|
|
prog = None
|
|
if parent_script:
|
|
prog = parent_script + " " + os.path.basename(__file__).split('.')[0]
|
|
|
|
# set up argument parser for the script
|
|
parser = argparse.ArgumentParser(
|
|
prog = prog,
|
|
description = "Create a batch of images based on a configuration file",
|
|
formatter_class = argparse.RawTextHelpFormatter,
|
|
epilog = """Environment Variables:
|
|
BULKGEN_COMFY_HOST / COMFY_HOST = Default comfyui host address.
|
|
BULKGEN_COMFY_PORT / COMFY_PORT = Default comfyui port number.
|
|
""")
|
|
parser.add_argument('--no-color' , action='store_true' , help="Disable colored output.")
|
|
parser.add_argument('--host' , type=str , help="ComfyUI host address.")
|
|
parser.add_argument('--port' , type=str , help="ComfyUI port.")
|
|
parser.add_argument('config_files', nargs='+', metavar='FILE', help="One or more configuration files to generate images from.")
|
|
args = parser.parse_args(args=args)
|
|
|
|
# if the user requested to disable colors, call disable_colors()
|
|
if args.no_color:
|
|
disable_colors()
|
|
|
|
# try to connect to ComfyUI,
|
|
# the host and port are determined by the --host and --port flags
|
|
# (if not specified, the values from environment variables are used)
|
|
comfy_host = args.host or DEFAULT_HOST
|
|
comfy_port = args.port or DEFAULT_PORT
|
|
comfyapi = ComfyAPI(comfy_host, comfy_port)
|
|
if not comfyapi.check_connection():
|
|
fatal_error(f"Could not connect to ComfyUI at \"{comfy_host}:{comfy_port}\"",
|
|
"Use the --host and --port flags to specify the host and port.",
|
|
"Alternatively, you can set the BULKGEN_COMFY_HOST and BULKGEN_COMFY_PORT environment variables.")
|
|
|
|
image_builder = ImageBuilder(comfyapi)
|
|
|
|
# loop through the configuration files and generate images
|
|
for config_file in args.config_files:
|
|
config = Config.from_file(config_file)
|
|
all_genparams = config.all_genparams
|
|
custom_styles = config.custom_styles
|
|
|
|
for prompt_index, genparams in enumerate(all_genparams):
|
|
filename_prefix = f"{FILE_PREFIX}_p{prompt_index:02}"
|
|
genparams_list = genparams.generate_combinations(name_prefix=filename_prefix)
|
|
for params in genparams_list:
|
|
image_bytes = image_builder(params)
|
|
with open(f"{params.filename}.png", "wb") as f:
|
|
f.write(image_bytes)
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |