From 3e049295150bdfe8973f820019443d2956d87506 Mon Sep 17 00:00:00 2001 From: aladac Date: Sun, 17 May 2026 18:15:33 +0200 Subject: [PATCH] feat(workflow): support UNet-only Flux.1 D checkpoints via DualCLIPLoader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- tensors/cli.py | 2 +- tensors/comfyui.py | 191 ++++++++++++++++++++++++++++++++++++++++++ tensors/config.py | 52 +++++++++++- tests/test_tensors.py | 165 ++++++++++++++++++++++++++++++++++++ 4 files changed, 408 insertions(+), 2 deletions(-) diff --git a/tensors/cli.py b/tensors/cli.py index 16228e7..03ad951 100644 --- a/tensors/cli.py +++ b/tensors/cli.py @@ -788,7 +788,7 @@ def generate( # noqa: PLR0915 help=( "Override detected model family " "(pony, illustrious, sdxl, sdxl_lightning, sdxl_turbo, " - "sd15, sd15_lcm, flux, flux_schnell, zimage)" + "sd15, sd15_lcm, flux, flux_schnell, flux_unet, zimage)" ), ), ] = None, diff --git a/tensors/comfyui.py b/tensors/comfyui.py index 531e029..9d33bb3 100644 --- a/tensors/comfyui.py +++ b/tensors/comfyui.py @@ -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 # Uses separate VAE loader for better quality with modern models DEFAULT_WORKFLOW_TEMPLATE: dict[str, Any] = { @@ -887,6 +975,87 @@ def _build_flux_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( prompt: str, negative_prompt: str = "", @@ -974,6 +1143,28 @@ def _build_workflow( 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) # Set seed (random if -1) diff --git a/tensors/config.py b/tensors/config.py index d1bd57f..f7fb876 100644 --- a/tensors/config.py +++ b/tensors/config.py @@ -671,6 +671,27 @@ MODEL_FAMILY_DEFAULTS: dict[str, dict[str, Any]] = { "steps": 4, "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": { "quality_prefix": "", "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 """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: 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() 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 # field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony" # but are architecturally Flux and need the Flux workflow). diff --git a/tests/test_tensors.py b/tests/test_tensors.py index 609edcd..3042970 100644 --- a/tests/test_tensors.py +++ b/tests/test_tensors.py @@ -320,6 +320,73 @@ class TestModelFamilyDetection: assert detect_model_family("model.safetensors", "Flux.1 D") == "flux" 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: """Test detecting SDXL family variants.""" from tensors.config import detect_model_family @@ -533,6 +600,104 @@ class TestFluxWorkflowBuilder: 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: """Tests for display formatting functions."""