feat(workflow): support UNet-only Flux.1 D checkpoints via DualCLIPLoader

Adds a `flux_unet` model family for Flux.1 D checkpoints that ship only the
UNet weights (no baked-in CLIP/T5/VAE) and require external encoder loading.

Affected checkpoints: lust_v10, cyberrealisticFlux_v25,
getphatFLUXReality_v11Softcore, moodyDesireMix_v20PRO,
fcFluxPonyPerfectBase. These live in `models/diffusion_models/` (or
`models/unet/`) on the ComfyUI host and are loaded via UNETLoader instead
of CheckpointLoaderSimple.

Detection:
- New `FLUX_UNET_ONLY_PATTERNS` constant in config.py lists known
  UNet-only filename substrings.
- `detect_model_family()` returns `flux_unet` when any pattern matches,
  taking precedence over base_model field (same architecture-override
  pattern used for FluxPony hybrids).

Workflow:
- New flux_unet workflow in comfyui.py uses UNETLoader + DualCLIPLoader
  (clip_l.safetensors + t5xxl_fp16.safetensors, type=flux) + VAELoader
  (ae.safetensors) wired into the same Flux KSampler graph as the
  monolithic flux family.
- Family defaults inherit from flux (sampler/scheduler/steps/cfg) with
  external_clip flag set.

Smoke-tested end-to-end on madcat with getphatFLUXReality_v11Softcore
producing valid output. Flux.2 Klein checkpoints (lust_v10,
moodyDesireMix) still fail because they require a different text encoder
(qwen_3_8b) — see follow-up commit.

13 new tests; 240 -> 253 total.
This commit is contained in:
2026-05-17 18:15:33 +02:00
parent af61ba79c9
commit 3e04929515
4 changed files with 408 additions and 2 deletions
+1 -1
View File
@@ -788,7 +788,7 @@ def generate( # noqa: PLR0915
help=( help=(
"Override detected model family " "Override detected model family "
"(pony, illustrious, sdxl, sdxl_lightning, sdxl_turbo, " "(pony, illustrious, sdxl, sdxl_lightning, sdxl_turbo, "
"sd15, sd15_lcm, flux, flux_schnell, zimage)" "sd15, sd15_lcm, flux, flux_schnell, flux_unet, zimage)"
), ),
), ),
] = None, ] = None,
+191
View File
@@ -742,6 +742,94 @@ FLUX_WORKFLOW_TEMPLATE: dict[str, Any] = {
} }
# Flux UNet-only workflow template.
#
# Used for checkpoints that ship without CLIP/T5/VAE baked in (e.g. Flux Dev's
# split-file release, lust_v10, cyberrealisticFlux, getphatFLUXReality,
# moodyDesireMix, fcFluxPony*). Structurally identical to FLUX_WORKFLOW_TEMPLATE
# but the single CheckpointLoaderSimple (node "100") is replaced by three
# separate loaders that share the same node IDs the rest of the graph already
# expects:
#
# "100" UNETLoader → outputs MODEL (slot 0)
# "101" DualCLIPLoader → outputs CLIP (slot 0)
# "102" VAELoader → outputs VAE (slot 0)
#
# Downstream nodes that referenced ["100", 1] (clip) now read ["101", 0];
# downstream nodes that referenced ["100", 2] (vae) now read ["102", 0].
FLUX_UNET_WORKFLOW_TEMPLATE: dict[str, Any] = {
"100": {
"class_type": "UNETLoader",
"inputs": {"unet_name": "", "weight_dtype": "default"},
},
"101": {
"class_type": "DualCLIPLoader",
"inputs": {
"clip_name1": "clip_l.safetensors",
"clip_name2": "t5xxl_fp16.safetensors",
"type": "flux",
},
},
"102": {
"class_type": "VAELoader",
"inputs": {"vae_name": "ae.safetensors"},
},
"120": {
"class_type": "ModelSamplingFlux",
"inputs": {
"model": ["100", 0],
"max_shift": 1.15,
"base_shift": 0.5,
"width": 1024,
"height": 1024,
},
},
"130": {
"class_type": "CLIPTextEncode",
"inputs": {"text": "", "clip": ["101", 0]},
},
"131": {
"class_type": "CLIPTextEncode",
"inputs": {"text": "", "clip": ["101", 0]},
},
"132": {
"class_type": "ConditioningZeroOut",
"inputs": {"conditioning": ["131", 0]},
},
"140": {
"class_type": "FluxGuidance",
"inputs": {"conditioning": ["130", 0], "guidance": 3.5},
},
"150": {
"class_type": "EmptySD3LatentImage",
"inputs": {"width": 1024, "height": 1024, "batch_size": 1},
},
"160": {
"class_type": "KSampler",
"inputs": {
"seed": 0,
"steps": 20,
"cfg": 1.0,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1.0,
"model": ["120", 0],
"positive": ["140", 0],
"negative": ["132", 0],
"latent_image": ["150", 0],
},
},
"170": {
"class_type": "VAEDecode",
"inputs": {"samples": ["160", 0], "vae": ["102", 0]},
},
"180": {
"class_type": "SaveImage",
"inputs": {"filename_prefix": "flux", "images": ["170", 0]},
},
}
# Default SDXL/Illustrious/Pony compatible workflow template # Default SDXL/Illustrious/Pony compatible workflow template
# Uses separate VAE loader for better quality with modern models # Uses separate VAE loader for better quality with modern models
DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = { DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = {
@@ -887,6 +975,87 @@ def _build_flux_workflow(
return workflow return workflow
def _build_flux_unet_workflow(
prompt: str,
model: str | None,
seed: int,
steps: int,
sampler: str,
scheduler: str,
width: int,
height: int,
batch_size: int,
lora_name: str | None,
lora_strength: float,
vae: str | None,
guidance: float,
clip_l: str,
clip_t5: str,
) -> dict[str, Any]:
"""Build a Flux workflow for UNet-only checkpoints (split CLIP/T5/VAE).
Structurally identical to ``_build_flux_workflow`` but uses UNETLoader +
DualCLIPLoader + VAELoader (nodes 100/101/102) instead of a single
CheckpointLoaderSimple. The rest of the graph (ModelSamplingFlux, KSampler,
FluxGuidance, etc.) is unchanged.
"""
workflow = copy.deepcopy(FLUX_UNET_WORKFLOW_TEMPLATE)
# Seed (random if -1)
actual_seed = seed if seed >= 0 else random.randint(0, 2**32 - 1)
# UNet checkpoint
if model:
workflow["100"]["inputs"]["unet_name"] = model
# CLIP/T5 (configurable so future variants can swap fp16 → fp8, etc.)
workflow["101"]["inputs"]["clip_name1"] = clip_l
workflow["101"]["inputs"]["clip_name2"] = clip_t5
# External VAE — fall back to ae.safetensors from the template if unset
if vae:
workflow["102"]["inputs"]["vae_name"] = vae
# ModelSamplingFlux must match the latent dimensions
workflow["120"]["inputs"]["width"] = width
workflow["120"]["inputs"]["height"] = height
# Prompts (positive only — negative is zero'd via ConditioningZeroOut)
workflow["130"]["inputs"]["text"] = prompt
# FluxGuidance carries the real prompt-adherence dial
workflow["140"]["inputs"]["guidance"] = guidance
# Latent
workflow["150"]["inputs"]["width"] = width
workflow["150"]["inputs"]["height"] = height
workflow["150"]["inputs"]["batch_size"] = batch_size
# KSampler — cfg stays 1.0
workflow["160"]["inputs"]["seed"] = actual_seed
workflow["160"]["inputs"]["steps"] = steps
workflow["160"]["inputs"]["sampler_name"] = sampler
workflow["160"]["inputs"]["scheduler"] = scheduler
# Optional LoRA injected between UNet/CLIP loaders and downstream consumers
if lora_name:
workflow["110"] = {
"class_type": "LoraLoader",
"inputs": {
"model": ["100", 0],
"clip": ["101", 0],
"lora_name": lora_name,
"strength_model": lora_strength,
"strength_clip": lora_strength,
},
}
workflow["120"]["inputs"]["model"] = ["110", 0]
workflow["130"]["inputs"]["clip"] = ["110", 1]
workflow["131"]["inputs"]["clip"] = ["110", 1]
return workflow
def _build_workflow( def _build_workflow(
prompt: str, prompt: str,
negative_prompt: str = "", negative_prompt: str = "",
@@ -974,6 +1143,28 @@ def _build_workflow(
guidance=_resolve_flux_guidance(guidance, cfg, defaults), guidance=_resolve_flux_guidance(guidance, cfg, defaults),
) )
# UNet-only Flux checkpoints (no baked-in CLIP/T5/VAE): use the split-loader
# variant. Triggered by family="flux_unet" — also covers any family whose
# preset opts in via external_clip=True.
if family == "flux_unet" or defaults.get("external_clip"):
return _build_flux_unet_workflow(
prompt=prompt,
model=model,
seed=seed,
steps=resolved_steps,
sampler=resolved_sampler,
scheduler=resolved_scheduler,
width=resolved_width,
height=resolved_height,
batch_size=batch_size,
lora_name=lora_name,
lora_strength=lora_strength,
vae=resolved_vae,
guidance=_resolve_flux_guidance(guidance, cfg, defaults),
clip_l=defaults.get("clip_l", "clip_l.safetensors"),
clip_t5=defaults.get("clip_t5", "t5xxl_fp16.safetensors"),
)
workflow = copy.deepcopy(DEFAULT_WORKFLOW_TEMPLATE) workflow = copy.deepcopy(DEFAULT_WORKFLOW_TEMPLATE)
# Set seed (random if -1) # Set seed (random if -1)
+51 -1
View File
@@ -671,6 +671,27 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
"steps": 4, "steps": 4,
"vae": "ae.safetensors", "vae": "ae.safetensors",
}, },
# UNet-only Flux checkpoints — same architecture as "flux" but the file
# ships without CLIP/T5/VAE baked in. Workflow must load them externally
# via UNETLoader + DualCLIPLoader + VAELoader instead of CheckpointLoaderSimple.
# external_clip=True signals this to the workflow builder.
"flux_unet": {
"quality_prefix": "",
"negative_prompt": "",
"width": 1024,
"height": 1024,
"portrait": (832, 1216),
"landscape": (1216, 832),
"cfg": 1.0,
"guidance": 3.5,
"sampler": "euler",
"scheduler": "simple",
"steps": 20,
"vae": "ae.safetensors",
"external_clip": True,
"clip_l": "clip_l.safetensors",
"clip_t5": "t5xxl_fp16.safetensors",
},
"zimage": { "zimage": {
"quality_prefix": "", "quality_prefix": "",
"negative_prompt": "", "negative_prompt": "",
@@ -687,6 +708,28 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = {
} }
# UNet-only Flux checkpoint filename substrings (case-insensitive). These ship
# without baked-in CLIP/T5/VAE, so they require UNETLoader + DualCLIPLoader +
# VAELoader instead of CheckpointLoaderSimple. Matched on the lowercased
# filename via simple substring containment.
#
# Add new patterns here as we encounter them — order doesn't matter, first
# match wins.
FLUX_UNET_ONLY_PATTERNS: tuple[str, ...] = (
"lust_", # lust_v10.safetensors (Flux.2 Klein 9B-base)
# Note: bare "lust" would falsely match "illustrious" — keep the underscore.
"cyberrealisticflux", # cyberrealisticFlux_v25.safetensors
"getphatflux", # getphatFLUXReality_v11Softcore.safetensors
"moodydesire", # moodyDesireMix_v20PRO.safetensors
"fcfluxpony", # fcFluxPonyPerfectBase_fcFluxPerfectBase.safetensors
)
def _is_flux_unet_only(name_lower: str) -> bool:
"""True if the lowercased filename matches a known UNet-only Flux pattern."""
return any(p in name_lower for p in FLUX_UNET_ONLY_PATTERNS)
def detect_model_family(model_name: str, base_model: str | None = None) -> str | None: # noqa: PLR0911 def detect_model_family(model_name: str, base_model: str | None = None) -> str | None: # noqa: PLR0911
"""Detect model family from filename or CivitAI base_model field. """Detect model family from filename or CivitAI base_model field.
@@ -696,11 +739,18 @@ def detect_model_family(model_name: str, base_model: str | None = None) -> str |
Returns: Returns:
Model family key (pony, illustrious, sdxl, sdxl_lightning, sdxl_turbo, Model family key (pony, illustrious, sdxl, sdxl_lightning, sdxl_turbo,
sd15, sd15_lcm, flux, flux_schnell, zimage) or None if unknown sd15, sd15_lcm, flux, flux_schnell, flux_unet, zimage) or None if unknown
""" """
name_lower = model_name.lower() name_lower = model_name.lower()
base_lower = (base_model or "").lower() base_lower = (base_model or "").lower()
# UNet-only Flux override: must run BEFORE the generic flux check below,
# since some patterns ("cyberrealisticflux", "getphatflux", "fcfluxpony")
# also contain the substring "flux". Filename wins over base_model
# field — these checkpoints are often mis-tagged on CivitAI.
if _is_flux_unet_only(name_lower):
return "flux_unet"
# Architecture override: filename containing "flux" wins over any base_model # Architecture override: filename containing "flux" wins over any base_model
# field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony" # field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony"
# but are architecturally Flux and need the Flux workflow). # but are architecturally Flux and need the Flux workflow).
+165
View File
@@ -320,6 +320,73 @@ class TestModelFamilyDetection:
assert detect_model_family("model.safetensors", "Flux.1 D") == "flux" assert detect_model_family("model.safetensors", "Flux.1 D") == "flux"
assert detect_model_family("model.safetensors", "Flux.1 S schnell") == "flux_schnell" assert detect_model_family("model.safetensors", "Flux.1 S schnell") == "flux_schnell"
def test_detect_flux_unet_lust(self) -> None:
"""lust_*.safetensors → flux_unet (no 'flux' in name, custom pattern)."""
from tensors.config import detect_model_family
assert detect_model_family("lust_v10.safetensors") == "flux_unet"
assert detect_model_family("LUST_v10.safetensors") == "flux_unet"
def test_detect_flux_unet_cyberrealistic(self) -> None:
"""cyberrealisticFlux_*.safetensors → flux_unet (intercepts generic 'flux' match)."""
from tensors.config import detect_model_family
assert detect_model_family("cyberrealisticFlux_v25.safetensors") == "flux_unet"
def test_detect_flux_unet_getphat(self) -> None:
"""getphatFLUXReality_*.safetensors → flux_unet."""
from tensors.config import detect_model_family
assert detect_model_family("getphatFLUXReality_v11Softcore.safetensors") == "flux_unet"
def test_detect_flux_unet_moody(self) -> None:
"""moodyDesireMix_*.safetensors → flux_unet (no 'flux' in name)."""
from tensors.config import detect_model_family
assert detect_model_family("moodyDesireMix_v20PRO.safetensors") == "flux_unet"
def test_detect_flux_unet_fcfluxpony(self) -> None:
"""fcFluxPony*.safetensors → flux_unet (intercepts flux + fluxpony)."""
from tensors.config import detect_model_family
assert (
detect_model_family("fcFluxPonyPerfectBase_fcFluxPerfectBase.safetensors")
== "flux_unet"
)
def test_detect_flux_unet_overrides_base_model(self) -> None:
"""Filename UNet-only pattern wins over a (likely wrong) CivitAI base_model tag."""
from tensors.config import detect_model_family
# Even if CivitAI claims "SDXL 1.0", the filename pattern wins.
assert detect_model_family("lust_v10.safetensors", "SDXL 1.0") == "flux_unet"
assert (
detect_model_family("cyberrealisticFlux_v25.safetensors", "Pony") == "flux_unet"
)
def test_flux_unet_family_defaults_has_external_clip(self) -> None:
"""flux_unet preset advertises external_clip + clip filenames."""
from tensors.config import MODEL_FAMILY_DEFAULTS
defaults = MODEL_FAMILY_DEFAULTS["flux_unet"]
assert defaults["external_clip"] is True
assert defaults["clip_l"] == "clip_l.safetensors"
assert defaults["clip_t5"] == "t5xxl_fp16.safetensors"
# Sanity: same sampling profile as flux
assert defaults["cfg"] == 1.0
assert defaults["guidance"] == 3.5
assert defaults["vae"] == "ae.safetensors"
def test_get_model_generation_defaults_flux_unet(self) -> None:
"""flux_unet model resolves to the flux_unet preset with external_clip set."""
from tensors.config import get_model_generation_defaults
defaults = get_model_generation_defaults("lust_v10.safetensors")
assert defaults["family"] == "flux_unet"
assert defaults["external_clip"] is True
assert defaults["sampler"] == "euler"
assert defaults["scheduler"] == "simple"
def test_detect_sdxl_variants(self) -> None: def test_detect_sdxl_variants(self) -> None:
"""Test detecting SDXL family variants.""" """Test detecting SDXL family variants."""
from tensors.config import detect_model_family from tensors.config import detect_model_family
@@ -533,6 +600,104 @@ class TestFluxWorkflowBuilder:
assert "120" not in wf assert "120" not in wf
class TestFluxUnetWorkflowBuilder:
"""Tests for the UNet-only Flux workflow (split CLIP/T5/VAE loaders)."""
def test_build_workflow_flux_unet_uses_dual_clip_loader(self) -> None:
"""flux_unet checkpoints emit UNETLoader + DualCLIPLoader + VAELoader and NO CheckpointLoaderSimple."""
from tensors.comfyui import _build_workflow
wf = _build_workflow(prompt="a cat", model="lust_v10.safetensors")
# Three split loaders at the canonical IDs
assert wf["100"]["class_type"] == "UNETLoader"
assert wf["101"]["class_type"] == "DualCLIPLoader"
assert wf["102"]["class_type"] == "VAELoader"
# The combined checkpoint loader must NOT appear anywhere.
for node in wf.values():
assert node["class_type"] != "CheckpointLoaderSimple"
# UNet filename plumbed through
assert wf["100"]["inputs"]["unet_name"] == "lust_v10.safetensors"
# DualCLIPLoader configured for flux with both encoders
clip_inputs = wf["101"]["inputs"]
assert clip_inputs["clip_name1"] == "clip_l.safetensors"
assert clip_inputs["clip_name2"] == "t5xxl_fp16.safetensors"
assert clip_inputs["type"] == "flux"
# VAE defaults to ae.safetensors
assert wf["102"]["inputs"]["vae_name"] == "ae.safetensors"
# Downstream wiring: CLIPTextEncode reads from DualCLIPLoader, VAEDecode from VAELoader
assert wf["130"]["inputs"]["clip"] == ["101", 0]
assert wf["131"]["inputs"]["clip"] == ["101", 0]
assert wf["170"]["inputs"]["vae"] == ["102", 0]
# ModelSamplingFlux still reads MODEL from node 100 (now UNETLoader)
assert wf["120"]["inputs"]["model"] == ["100", 0]
def test_flux_unet_inherits_flux_sampling_profile(self) -> None:
"""flux_unet locks KSampler.cfg to 1.0 and exposes the FluxGuidance dial."""
from tensors.comfyui import _build_workflow
wf = _build_workflow(
prompt="a cat", model="lust_v10.safetensors", cfg=7.5
)
assert wf["160"]["inputs"]["cfg"] == 1.0
# The caller's cfg=7.5 should re-route to FluxGuidance (same precedence as plain flux)
assert wf["140"]["inputs"]["guidance"] == 7.5
def test_flux_unet_lora_injection_wires_split_loaders(self) -> None:
"""LoRA on flux_unet injects node 110 reading MODEL from UNETLoader and CLIP from DualCLIPLoader."""
from tensors.comfyui import _build_workflow
wf = _build_workflow(
prompt="a cat",
model="lust_v10.safetensors",
lora_name="my_style.safetensors",
lora_strength=0.6,
)
assert wf["110"]["class_type"] == "LoraLoader"
assert wf["110"]["inputs"]["lora_name"] == "my_style.safetensors"
assert wf["110"]["inputs"]["strength_model"] == 0.6
# LoRA reads model from UNETLoader, clip from DualCLIPLoader
assert wf["110"]["inputs"]["model"] == ["100", 0]
assert wf["110"]["inputs"]["clip"] == ["101", 0]
# Downstream consumers now read from the LoRA outputs
assert wf["120"]["inputs"]["model"] == ["110", 0]
assert wf["130"]["inputs"]["clip"] == ["110", 1]
assert wf["131"]["inputs"]["clip"] == ["110", 1]
def test_flux_unet_external_vae_overrides_default(self) -> None:
"""Caller-provided VAE replaces ae.safetensors on the VAELoader (no new node)."""
from tensors.comfyui import _build_workflow
wf = _build_workflow(
prompt="a cat",
model="lust_v10.safetensors",
vae="other_vae.safetensors",
)
assert wf["102"]["inputs"]["vae_name"] == "other_vae.safetensors"
# And VAEDecode still wires through node 102 — no shadow node 171.
assert wf["170"]["inputs"]["vae"] == ["102", 0]
assert "171" not in wf
def test_flux_unet_via_explicit_family_override(self) -> None:
"""A non-pattern filename still gets the UNet workflow when -F flux_unet is forced.
We can't pass --family directly to _build_workflow (it auto-detects from
the filename), but a checkpoint matching a UNet-only pattern proves the
family→workflow dispatch end-to-end.
"""
from tensors.comfyui import _build_workflow
# moodyDesireMix has no "flux" in name but must route to the UNet workflow.
wf = _build_workflow(prompt="a cat", model="moodyDesireMix_v20PRO.safetensors")
assert wf["100"]["class_type"] == "UNETLoader"
assert wf["101"]["class_type"] == "DualCLIPLoader"
class TestDisplayFormatters: class TestDisplayFormatters:
"""Tests for display formatting functions.""" """Tests for display formatting functions."""