nodes/JeyzerMC-comfyui-flakes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
85f030a874780a4bd3f58ea0b25ca4c4020259ce
JeyzerMC-comfyui-flakes/full_flake_node.py
T

558 lines
23 KiB
Python
Raw Blame History

from __future__ import annotations
import json
import logging
import os
from datetime import datetime
import folder_paths
import comfy.sd
import comfy.utils
from nodes import CLIPTextEncode, EmptyLatentImage, ControlNetApplyAdvanced, ControlNetLoader, LoraLoader
from . import flake_io
from .flake_io import _resolve_model_name
def _split_stem_for_filename(stem: str) -> tuple[str, str, str, str]:
"""Parse one filename-prefix stem into placement buckets, honoring the
``**`` suffix marker (#290).
The last path segment is the filename contribution; preceding segments are
folders. A ``**`` before the folder portion makes that folder a *suffix*
(placed after other flakes' folders); a ``**`` before the file portion
makes that name a *suffix* (placed after other flakes' names).
Returns ``(folder_prefix, folder_suffix, file_prefix, file_suffix)`` with
each element possibly empty. Folder values keep their trailing ``/``.
"""
folder_prefix = folder_suffix = ""
file_prefix = file_suffix = ""
slash = stem.rfind("/")
if slash >= 0:
folder_region = stem[:slash + 1]
file_region = stem[slash + 1:]
else:
folder_region = ""
file_region = stem
if folder_region:
if folder_region.startswith("**"):
folder_suffix = folder_region[2:].replace("**", "")
else:
folder_prefix = folder_region.replace("**", "")
if file_region:
before, sep, after = file_region.partition("**")
file_prefix = before.replace("**", "")
if sep:
file_suffix = after.replace("**", "")
return folder_prefix, folder_suffix, file_prefix, file_suffix
def _build_filename_prefix(preset_name: str, stems: list[str]) -> str:
"""Build a filename prefix from preset name and flake stems.
Folders accumulate as ``<folder_prefixes><folder_suffixes>`` and the
filename as ``<time>_<file_prefixes>_<file_suffixes>``, where prefixes come
from each stem in stack order and suffixes (marked with ``**``) follow all
prefixes — see :func:`_split_stem_for_filename` (#290).
"""
folder_prefixes: list[str] = []
folder_suffixes: list[str] = []
file_prefixes: list[str] = []
file_suffixes: list[str] = []
for stem in stems:
if not stem:
continue
fpre, fsuf, npre, nsuf = _split_stem_for_filename(stem)
if fpre:
folder_prefixes.append(fpre)
if fsuf:
folder_suffixes.append(fsuf)
if npre:
file_prefixes.append(npre)
if nsuf:
file_suffixes.append(nsuf)
path = (preset_name + "/") if preset_name else ""
path += "".join(folder_prefixes) + "".join(folder_suffixes)
now = datetime.now()
path += now.strftime("%y%m%d") + "/"
filename = now.strftime("%H%M%S")
name_parts = file_prefixes + file_suffixes
if name_parts:
filename += "_" + "_".join(name_parts)
return path + filename
def _carry_controlnets(src_cond, dst_cond):
"""Copy the controlnet chain from upstream conditioning onto freshly
re-encoded conditioning (#257).
Flake nodes re-encode the accumulated prompt text every time, producing
fresh conditioning that loses the ``control`` key — which is how ComfyUI
chains ControlNets (see ControlNetApplyAdvanced). Without this, a
controlnet applied by an upstream Flake node is silently dropped and only
the last node's controlnets survive. Carrying ``control`` forward lets this
node's own controlnets chain on top via ``set_previous_controlnet``.
"""
if not src_cond or not dst_cond:
return dst_cond
src_meta = src_cond[0][1] if len(src_cond[0]) > 1 else {}
control = src_meta.get("control")
if control is None:
return dst_cond
out = []
for t in dst_cond:
d = t[1].copy()
d["control"] = control
if "control_apply_to_uncond" in src_meta:
d["control_apply_to_uncond"] = src_meta["control_apply_to_uncond"]
out.append([t[0], d])
return out
def _resolve_lora_name(stem_or_name: str) -> str:
result = _resolve_model_name("loras", stem_or_name)
available = folder_paths.get_filename_list("loras")
if result.replace("\\", "/") in {p.replace("\\", "/") for p in available}:
return result
raise FileNotFoundError(f"LoRA '{stem_or_name}' not found in models/loras/")
def _load_preset_bundle(preset_name: str, model_family: str | None = None):
"""Load a model preset and return (model_bundle, generation_data, sampling_preset).
``model_family`` is used to derive the ``img/<family>/`` output folder
prefix (#240). It's optional for backwards compatibility but should be
passed by the node so the structural prefix is present.
"""
preset_data = flake_io.load_preset(preset_name)
# --- Load checkpoint ----------------------------------------------------
ckpt_name = _resolve_model_name("checkpoints", preset_data.checkpoint)
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
if not ckpt_path or not os.path.isfile(ckpt_path):
raise FileNotFoundError(
f"Checkpoint '{preset_data.checkpoint}' not found in models/checkpoints/"
)
embedding_dir = folder_paths.get_folder_paths("embeddings")
model, clip, vae, _ = comfy.sd.load_checkpoint_guess_config(
ckpt_path,
output_vae=True,
output_clip=True,
embedding_directory=embedding_dir,
)
# --- Clip skip ----------------------------------------------------------
if preset_data.clip_skip:
clip = clip.clone()
clip.clip_layer(preset_data.clip_skip)
# --- Optional VAE override ----------------------------------------------
if preset_data.vae:
vae_name = _resolve_model_name("vae", preset_data.vae)
vae_path = folder_paths.get_full_path("vae", vae_name)
if vae_path and os.path.isfile(vae_path):
vae_sd = comfy.utils.load_torch_file(vae_path)
vae = comfy.sd.VAE(sd=vae_sd)
# --- Optional text encoder override ------------------------------------
if preset_data.text_encoder:
te_name = _resolve_model_name("text_encoders", preset_data.text_encoder)
te_path = folder_paths.get_full_path("text_encoders", te_name)
if te_path and os.path.isfile(te_path):
te_sd = comfy.utils.load_torch_file(te_path)
_, clip, _ = comfy.sd.load_checkpoint_guess_config(
te_path,
output_vae=False,
output_clip=True,
embedding_directory=embedding_dir,
)
if preset_data.clip_skip:
clip = clip.clone()
clip.clip_layer(preset_data.clip_skip)
# --- Encode prompts -----------------------------------------------------
encoder = CLIPTextEncode()
pos_text = preset_data.positive.strip()
neg_text = preset_data.negative.strip()
positive = encoder.encode(clip, pos_text)[0] if pos_text else encoder.encode(clip, "")[0]
negative = encoder.encode(clip, neg_text)[0] if neg_text else encoder.encode(clip, "")[0]
# --- Latent -------------------------------------------------------------
width, height = preset_data.width, preset_data.height
latent = EmptyLatentImage().generate(width, height, 1)[0]
model_bundle = (model, clip, vae)
# Build the structural prefix `img/<family>/` (#240) plus the explicit
# filename_prefix field. We do NOT inherit from the preset's yaml location
# on disk (#217). The img/<family>/ segment is the project's output folder
# convention; without it, outputs land at output/<rest>/... instead of
# output/img/<family>/<rest>/...
family_folder = flake_io._family_folder(model_family) if model_family else None
structural = f"img/{family_folder}/" if family_folder else ""
prefix_stem = (preset_data.filename_prefix or "").strip()
stems = []
if structural:
# stems entries containing "/" land in folder_parts (see _build_filename_prefix).
stems.append(structural)
if prefix_stem:
stems.append(prefix_stem)
filename_state = {
"preset": "",
"stems": stems,
"checkpoint": preset_data.checkpoint,
"vae": preset_data.vae or "baked-in",
"text_encoder": preset_data.text_encoder or "baked-in",
"loras": [],
}
generation_data = (positive, negative, latent, width, height, pos_text, neg_text, filename_state)
sampling_preset = (preset_data.steps, preset_data.cfg, preset_data.sampler, preset_data.scheduler)
return model_bundle, generation_data, sampling_preset
def _apply_preset_overrides(overrides_json, model_bundle, generation_data, sampling_preset):
"""Overlay per-instance overrides (steps/cfg/sampler/scheduler/
filename_prefix) onto a loaded preset bundle (#279).
Empty/blank fields are ignored so the preset's own values pass through.
``filename_prefix`` is appended to the stem chain rather than replacing it.
"""
if not overrides_json:
return model_bundle, generation_data, sampling_preset
try:
ovr = json.loads(overrides_json) or {}
except (ValueError, TypeError):
return model_bundle, generation_data, sampling_preset
if not isinstance(ovr, dict) or not ovr:
return model_bundle, generation_data, sampling_preset
steps, cfg, sampler, scheduler = sampling_preset
def _conv(key, conv, cur):
v = ovr.get(key)
if v is None or v == "":
return cur
try:
return conv(v)
except (ValueError, TypeError):
return cur
steps = _conv("steps", int, steps)
cfg = _conv("cfg", float, cfg)
if ovr.get("sampler"):
sampler = str(ovr["sampler"])
if ovr.get("scheduler"):
scheduler = str(ovr["scheduler"])
sampling_preset = (steps, cfg, sampler, scheduler)
fp = (ovr.get("filename_prefix") or "").strip()
if fp and isinstance(generation_data, tuple) and len(generation_data) > 7 and isinstance(generation_data[7], dict):
fs = dict(generation_data[7])
fs["stems"] = list(fs.get("stems", [])) + [fp]
generation_data = generation_data[:7] + (fs,)
return model_bundle, generation_data, sampling_preset
_MODEL_FAMILIES = ["SDXL/Base", "SDXL/Illustrious", "SDXL/Pony", "ZImage/Base", "ZImage/Turbo", "Anima/Base", "Flux/Klein"]
class FlakeModelPreset:
@classmethod
def INPUT_TYPES(cls):
try:
preset_names = flake_io.list_presets()
except Exception:
preset_names = []
if not preset_names:
preset_names = ["No model preset is selected"]
return {
"required": {
"model_family": (_MODEL_FAMILIES, {"default": "SDXL/Base"}),
"preset": (["Select a preset..."] + preset_names,),
},
"optional": {
# Per-instance overrides (steps/cfg/sampler/scheduler/
# filename_prefix) as JSON. Managed by the node's widget; lets a
# preset be A/B-tested without editing its file (#279).
"overrides_json": ("STRING", {
"default": "",
"multiline": True,
"tooltip": "JSON override of steps/cfg/sampler/scheduler/filename_prefix. Managed by the Flake widget.",
}),
},
}
RETURN_TYPES = ("FLAKE_DATA",)
RETURN_NAMES = ("flake_data",)
FUNCTION = "execute"
CATEGORY = "flakes"
DESCRIPTION = (
"Load a model preset (checkpoint, VAE, prompts, resolution, sampler settings). "
"Outputs a single bundled flake_data pin for wiring into FlakeStack / FlakeCombo nodes."
)
@classmethod
def IS_CHANGED(cls, model_family: str, preset: str, overrides_json: str = ""):
# Re-run whenever the preset file's contents change on disk, even if the
# selected preset name (the node input) stays the same. Without this,
# ComfyUI caches the output and edits to Steps/CFG/etc. are ignored.
# The override JSON is part of the signature so override edits re-run too.
preset_name = preset.strip() if preset else ""
if not preset_name or preset_name in ("Select a preset...", "No model preset is selected"):
return f"{preset_name}:{overrides_json}"
try:
path = flake_io._resolve_preset_file(preset_name)
st = os.stat(path)
return f"{path}:{st.st_mtime_ns}:{st.st_size}:{overrides_json}"
except Exception:
return float("nan")
def execute(self, model_family: str, preset: str, overrides_json: str = ""):
preset_name = preset.strip() if preset else ""
if not preset_name or preset_name in ("Select a preset...", "No model preset is selected"):
raise ValueError("No model preset is selected — pick one from the dropdown.")
model_bundle, generation_data, sampling_preset = _load_preset_bundle(preset_name, model_family)
model_bundle, generation_data, sampling_preset = _apply_preset_overrides(
overrides_json, model_bundle, generation_data, sampling_preset
)
return ((model_bundle, generation_data, sampling_preset),)
class FlakeStack:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_family": (_MODEL_FAMILIES, {"default": "SDXL/Base"}),
"flake_data": ("FLAKE_DATA",),
"flakes_json": ("STRING", {
"multiline": True,
"default": "[]",
"tooltip": "JSON-encoded list of flake entries. Managed by the Full Flakes widget.",
}),
},
}
RETURN_TYPES = ("FLAKE_DATA",)
RETURN_NAMES = ("flake_data",)
FUNCTION = "execute"
CATEGORY = "flakes"
DESCRIPTION = (
"Compose flakes on top of an incoming flake_data bundle. "
"Applies LoRAs, merges prompts, resolution overrides and ControlNets. "
"Outputs updated flake_data for chaining into downstream nodes."
)
def execute(self, model_family: str, flake_data, flakes_json: str):
model_bundle, generation_data, sampling_preset = flake_data
model, clip, vae = model_bundle
positive_cond, negative_cond, latent, width, height, pos_text, neg_text = generation_data[:7]
steps, cfg, sampler, scheduler = sampling_preset
# --- Filename prefix state ----------------------------------------------
# Copy the incoming filename_state so we don't mutate the original tuple.
# Preserve existing stems from upstream FlakeStack nodes so that stems
# accumulate correctly across a chain of stacked nodes.
if len(generation_data) > 7 and isinstance(generation_data[7], dict):
filename_state = dict(generation_data[7])
filename_state["stems"] = list(filename_state.get("stems", []))
else:
filename_state = {"preset": "", "stems": []}
# --- Parse flakes_json --------------------------------------------------
try:
entries = json.loads(flakes_json) if flakes_json else []
except json.JSONDecodeError as exc:
raise ValueError(f"flakes_json is not valid JSON: {exc}") from exc
if not isinstance(entries, list):
raise ValueError("flakes_json must be a JSON list")
normalized: list[dict] = []
for i, entry in enumerate(entries):
if not isinstance(entry, dict):
logging.warning("[FlakeStack] skipping non-object entry %d: %r", i, entry)
continue
if entry.get("inline") or entry.get("name"):
normalized.append(entry)
else:
logging.warning("[FlakeStack] skipping entry %d (needs 'name' or 'inline')", i)
# --- Resolve flakes (skip bypassed) --------------------------------------
active_entries = [e for i, e in enumerate(normalized) if not e.get("bypassed")]
flakes = [flake_io.resolve(e) for e in active_entries]
# --- Apply LoRAs --------------------------------------------------------
lora_loader = LoraLoader()
lora_entries = []
for f in flakes:
for lr in f.loras:
if not lr.path:
continue
lora_name = _resolve_lora_name(lr.path)
model, clip = lora_loader.load_lora(model, clip, lora_name, lr.strength, lr.strength)
lora_entries.append({"name": lr.name or lr.path, "path": lr.path, "strength": lr.strength})
# Legacy single LoRA fallback
if f.lora_path:
lora_name = _resolve_lora_name(f.lora_path)
model, clip = lora_loader.load_lora(model, clip, lora_name, f.strength, f.strength)
lora_entries.append({"name": f.lora_path, "path": f.lora_path, "strength": f.strength})
# --- Build prompt text (existing → flakes) ------------------------------
pos_parts: list[str] = []
neg_parts: list[str] = []
if pos_text and pos_text.strip():
pos_parts.append(pos_text.strip())
if neg_text and neg_text.strip():
neg_parts.append(neg_text.strip())
for f in flakes:
if f.positive and f.positive.strip():
pos_parts.append(f.positive.strip())
if f.negative and f.negative.strip():
neg_parts.append(f.negative.strip())
new_pos_text = " BREAK ".join(pos_parts) if pos_parts else ""
new_neg_text = ", ".join(neg_parts) if neg_parts else ""
# --- Re-encode prompts --------------------------------------------------
# Re-encoding produces fresh conditioning that drops the ControlNet
# chain attached by upstream Flake nodes, so snapshot the incoming
# conditioning and carry its `control` forward afterwards (#257).
prev_positive_cond, prev_negative_cond = positive_cond, negative_cond
encoder = CLIPTextEncode()
positive_cond = encoder.encode(clip, new_pos_text)[0] if new_pos_text else encoder.encode(clip, "")[0]
negative_cond = encoder.encode(clip, new_neg_text)[0] if new_neg_text else encoder.encode(clip, "")[0]
positive_cond = _carry_controlnets(prev_positive_cond, positive_cond)
negative_cond = _carry_controlnets(prev_negative_cond, negative_cond)
# --- Target resolution (skip bypassed) ----------------------------------
# Computed before ControlNets so hint images can be fitted to the frame.
new_width, new_height = width, height
for f in flakes:
if f.resolution is not None:
new_width, new_height = f.resolution
break
# --- Apply ControlNets (from flakes) ------------------------------------
from .flake_compose import _load_cn_image, _fit_cn_image
# ControlNet "use image dimensions as resolution" override (#260): the
# first active CN entry flagged resolution_from_image drives the
# generation size, taking precedence over the flake resolution above.
for f in flakes:
hit = False
for cn in f.controlnets:
if cn.strength == 0 or not cn.resolution_from_image or not cn.image_name.strip():
continue
try:
ov = _load_cn_image(cn.image_name)
except Exception as exc:
logging.warning("[flakes] resolution_from_image failed for %s: %s", cn.image_name, exc)
continue
new_width, new_height = int(ov.shape[2]), int(ov.shape[1])
logging.info("[flakes] resolution override from CN image %s -> %dx%d", cn.image_name, new_width, new_height)
hit = True
break
if hit:
break
family_folder = flake_io._FAMILY_MAP.get(model_family, "sdxl")
cn_model_cache = {}
cn_loader = ControlNetLoader()
cn_apply = ControlNetApplyAdvanced()
for f in flakes:
for cn in f.controlnets:
if cn.strength == 0:
continue
model_name = cn.model_name
if cn.type:
inferred = flake_io.infer_cn_model(cn.type, family_folder)
if inferred:
model_name = inferred
if not model_name.strip():
print(f"[flakes] skipping controlnet entry with empty model_name (type={cn.type})")
continue
if not cn.image_name.strip():
print(f"[flakes] skipping controlnet entry with empty image_name (type={cn.type})")
continue
cn_resolved = flake_io.resolve_cn_model_name(model_name, family_folder)
if cn_resolved not in cn_model_cache:
cn_model_cache[cn_resolved] = cn_loader.load_controlnet(cn_resolved)[0]
cn_model = cn_model_cache[cn_resolved]
image = _load_cn_image(cn.image_name)
image = _fit_cn_image(image, new_width, new_height)
positive_cond, negative_cond = cn_apply.apply_controlnet(
positive_cond, negative_cond, cn_model, image,
cn.strength, cn.start_percent, cn.end_percent,
)
logging.info(
"[flakes] applied CN model=%s image=%s fit=%dx%d strength=%.2f start=%.2f end=%.2f",
cn_resolved, cn.image_name, new_width, new_height,
cn.strength, cn.start_percent, cn.end_percent,
)
# --- Latent (resolution resolved above, before ControlNets) -------------
if (new_width, new_height) != (width, height):
latent = EmptyLatentImage().generate(new_width, new_height, 1)[0]
# --- Lora metadata for Preview Flake Data ---------------------------------
filename_state["loras"] = lora_entries
# --- Filename prefix (skip bypassed) ------------------------------------
for f in flakes:
if f.output_stem:
filename_state["stems"].append(f.output_stem)
logging.info(
"[FlakeStack] %dx%s steps=%s cfg=%s flakes=%d",
new_width, new_height, steps, cfg, len(flakes),
)
out_model_bundle = (model, clip, vae)
out_generation_data = (positive_cond, negative_cond, latent, new_width, new_height, new_pos_text, new_neg_text, filename_state)
out_sampling_preset = (steps, cfg, sampler, scheduler)
return (
(out_model_bundle, out_generation_data, out_sampling_preset),
)
class FlakeCombo(FlakeStack):
"""Mutually-exclusive flake stack.
The backend behavior is identical to FlakeStack — it receives a
``flakes_json`` list and applies whatever entries are inside.
The frontend widget ensures that ``flakes_json`` contains exactly
one active flake at a time, and the JavaScript extension hooks
``app.queuePrompt`` to queue every combination across all
FlakeCombo / FlakeModelCombo nodes in the graph.
"""
pass
class FlakeModelCombo(FlakeModelPreset):
"""Multiple model preset selector.
The backend behavior is identical to FlakeModelPreset — it loads
a single preset from the ``preset`` dropdown. The frontend widget
manages a list of selected presets in ``node.properties`` and the
JavaScript extension cycles the ``preset`` widget value for each
combination when queueing.
"""
pass
Reference in New Issue View Git Blame Copy Permalink