From 12395a103fd0ca499c631dc6f460a8778e11490d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 15 Oct 2026 11:18:14 -0600 Subject: [PATCH] Load comfy files if they exist in flux.3 image --- .../diffusion_models/flux3/flux3.py | 127 ++++++++++-------- .../diffusion_models/flux3/src/checkpoints.py | 33 +++++ .../diffusion_models/flux3/src/transformer.py | 35 +++++ toolkit/models/v2/_mixin.py | 26 +++- 4 files changed, 164 insertions(+), 57 deletions(-) create mode 100644 extensions_built_in/diffusion_models/flux3/src/checkpoints.py diff --git a/extensions_built_in/diffusion_models/flux3/flux3.py b/extensions_built_in/diffusion_models/flux3/flux3.py index 71efd16..4bddee0 100644 --- a/extensions_built_in/diffusion_models/flux3/flux3.py +++ b/extensions_built_in/diffusion_models/flux3/flux3.py @@ -10,20 +10,23 @@ t2i and i2i (reference images as clean conditioning tokens). in an ``AdvancedPromptEmbeds`` (padded to a multiple of 80 tokens -- the DiT has no text mask, pad tokens are part of the trained context). -``model.name_or_path`` is the model package (hub repo id, default -``black-forest-labs/flux-3-image-klein-base``, or a local dir): +``model.name_or_path`` defaults to the vendor package +``black-forest-labs/flux-3-image-klein-base`` (or a local dir): config.json model.safetensors ae.safetensors text_encoder/ -Only the DiT and the image VAE come from the package, and they live in the -ComfyUI layout under MODELS_PATH: ``diffusion_models/flux_3_image_klein_base.safetensors`` -and ``vae/flux_3_vae.safetensors``. A file already there (or anywhere under its -category folder) is used in place; only a missing one is downloaded from the -package repo and renamed to that path. The package's -text_encoder/ and tokenizer/ are byte-identical to Qwen/Qwen3-VL-4B-Instruct, -so the text encoder loads from the Comfy-Org single file under -MODELS_PATH/text_encoders/ (downloaded there once when missing) with config -and tokenizer from the public Qwen repo. +The DiT and VAE live in the ComfyUI layout under MODELS_PATH, under the names +of BFL's ComfyUI repack ``black-forest-labs/flux-3-image-klein-comfy`` +(``src/checkpoints.py``): a repack file already there (the best precision +variant for the requested qtype -- int8_convrot attaches straight onto the +convrot8 backend; bf16; nvfp4 -- or anywhere under its category folder) is +used in place; only a missing one is downloaded from the package and renamed +to that spot (the package files are byte-identical to the repack's bf16 +ones). Naming the repack itself as ``name_or_path`` downloads its files +instead. The package's text_encoder/ and tokenizer/ are byte-identical to +Qwen/Qwen3-VL-4B-Instruct, so the text encoder loads from the Comfy-Org +single file under MODELS_PATH/text_encoders/ (downloaded there once when +missing) with config and tokenizer from the public Qwen repo. Flow convention matches ai-toolkit exactly (t=1 noise -> t=0 clean, target = noise - clean), so no time flips anywhere. @@ -47,6 +50,7 @@ from toolkit.basic import flush from toolkit.config_modules import GenerateImageConfig, ModelConfig from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel +from toolkit.models.v2.resolver import resolve_comfy_file from toolkit.paths import MODELS_PATH from toolkit.unloader import FakeTextEncoder from toolkit.samplers.custom_flowmatch_sampler import ( @@ -61,6 +65,7 @@ from .src.pipeline import ( pack_i2i_image_condition, pack_image_latents, ) +from .src.checkpoints import BASE_REPO, COMFY_REPO, PACKAGE_TO_COMFY from .src.text_encoder import ( DEFAULT_MAX_TOKENS, DEFAULT_PAD_MULTIPLE, @@ -74,12 +79,7 @@ if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO -DEFAULT_REPO = "black-forest-labs/flux-3-image-klein-base" -# package file -> ComfyUI-layout location under MODELS_PATH -COMFY_FILES = { - "model.safetensors": "diffusion_models/flux_3_image_klein_base.safetensors", - "ae.safetensors": "vae/flux_3_vae.safetensors", -} +DEFAULT_REPO = BASE_REPO # Static flow-matching timeshift used for the training timestep/noise # schedule (the package's sampler alpha). @@ -161,8 +161,12 @@ class Flux3Model(BaseModel): return None if os.path.isfile(base) and filename == "model.safetensors": return base # name_or_path points straight at the DiT file - if filename in COMFY_FILES: + if filename in PACKAGE_TO_COMFY: return self._comfy_weight(filename, base) + if base == COMFY_REPO: + if required: + raise FileNotFoundError(f"{COMFY_REPO} carries no {filename}") + return None from huggingface_hub import hf_hub_download try: @@ -176,43 +180,47 @@ class Flux3Model(BaseModel): return None def _comfy_weight(self, filename: str, repo_id: str) -> str: - """``MODELS_PATH//`` (or the local name anywhere - under the category folder); a miss downloads the package file and - renames it to that spot.""" - from toolkit.models.v2.resolver import find_file_recursive - - rel_path = COMFY_FILES[filename] - target = os.path.join(MODELS_PATH, rel_path) - if os.path.exists(target): - return target - found = find_file_recursive( - os.path.join(MODELS_PATH, os.path.dirname(rel_path)), - os.path.basename(rel_path), + """The comfy-layout file for a package weight: ``MODELS_PATH/`` + (or its name anywhere under the category folder) wins; a miss + downloads the repack's file there, or, for the vendor package, + downloads the package file and renames it to the same spot.""" + rel_path = PACKAGE_TO_COMFY[filename] + found = resolve_comfy_file( + rel_path, + COMFY_REPO, + hf_token=os.getenv("HF_TOKEN", None), + status_fn=self.print_and_status_update, + local_only=repo_id != COMFY_REPO, ) if found is not None: return found from huggingface_hub import hf_hub_download + target = os.path.join(MODELS_PATH, rel_path) self.print_and_status_update( f"Downloading {filename} from {repo_id} to {target}" ) - category_dir = os.path.dirname(target) - os.makedirs(category_dir, exist_ok=True) + # staged in a private subfolder (same filesystem, so the rename is a + # move): the package names collide with other models' comfy files + # (vae/ae.safetensors is the FLUX.1 VAE) and must never land there + staging = os.path.join(os.path.dirname(target), ".flux3_download") + os.makedirs(staging, exist_ok=True) path = hf_hub_download( repo_id=repo_id, filename=filename, token=os.getenv("HF_TOKEN", None), - local_dir=category_dir, + local_dir=staging, ) os.replace(path, target) - # the hub leaves download metadata under /.cache/huggingface - shutil.rmtree(os.path.join(category_dir, ".cache"), ignore_errors=True) + shutil.rmtree(staging, ignore_errors=True) return target def _load_package_config(self) -> Optional[dict]: if not self._package_config_checked: self._package_config_checked = True - path = self._package_file("config.json", required=False) + path = None + if self.model_config.name_or_path != COMFY_REPO: # the repack has none + path = self._package_file("config.json", required=False) if path is not None: with open(path, "r") as f: self.package_config = json.load(f) @@ -247,7 +255,6 @@ class Flux3Model(BaseModel): def _load_transformer(self) -> Flux3Transformer: dtype = self.torch_dtype self.print_and_status_update("Loading transformer (Flux3Transformer)") - dit_path = self._package_file("model.safetensors") params = self._transformer_params() cfg = self._load_package_config() if cfg is not None and cfg.get("quantization", None): @@ -262,25 +269,27 @@ class Flux3Model(BaseModel): "use the base package." ) - state_dict = load_file(dit_path) - state_dict = { - k: (v.to(dtype) if v.is_floating_point() else v) - for k, v in state_dict.items() - } - with torch.device("meta"): - transformer = Flux3Transformer(params) - needed = set(transformer.state_dict().keys()) - missing = sorted(needed - set(state_dict.keys())) - if missing: - raise ValueError(f"FLUX.3 checkpoint missing required keys: {missing[:8]}") - result = transformer.load_state_dict(state_dict, strict=False, assign=True) - if result.unexpected_keys: - self.print_and_status_update( - f" - ignoring {len(result.unexpected_keys)} unexpected checkpoint " - f"keys (e.g. {result.unexpected_keys[:3]})" + qtype = self.model_config.qtype if self.model_config.quantize else None + base = self.model_config.name_or_path + source = None + if base in Flux3Transformer.aitk_comfy_weight_names: + source = base # the repack named outright: downloaded when missing + elif base == BASE_REPO: + # a repack file already on disk (best precision variant for the + # qtype) wins over the package's own DiT; nothing is downloaded + # from the repack + source = Flux3Transformer.resolve_comfy_weights( + COMFY_REPO, qtype=qtype, local_only=True ) - del state_dict - flush() + if source is None: + source = self._package_file("model.safetensors") + transformer = Flux3Transformer.load_model( + source, + dtype=dtype, + config=params, + qtype=qtype, + quantize_on_load=False, # qtype ranks the comfy candidates only + ) attn_backend = self.model_config.model_kwargs.get("attention_backend", None) if attn_backend is not None: @@ -314,8 +323,14 @@ class Flux3Model(BaseModel): def _load_vae(self) -> Flux3ImageVAE: self.print_and_status_update("Loading image VAE") + from toolkit.models.v2.vae.flux2_kl import convert_diffusers_state_dict + vae = Flux3ImageVAE() - vae.load_state_dict(load_file(self._package_file("ae.safetensors"))) + state_dict = load_file(self._package_file("ae.safetensors")) + if any(k.startswith("encoder.down_blocks.") for k in state_dict): + # Comfy-Org's FLUX.2 VAE file: same weights, diffusers key layout + state_dict = convert_diffusers_state_dict(state_dict) + vae.load_state_dict(state_dict) vae.to(self.vae_torch_dtype).eval().requires_grad_(False) return vae diff --git a/extensions_built_in/diffusion_models/flux3/src/checkpoints.py b/extensions_built_in/diffusion_models/flux3/src/checkpoints.py new file mode 100644 index 0000000..09ab186 --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/checkpoints.py @@ -0,0 +1,33 @@ +"""Where FLUX.3 image-klein weights come from. + +BASE_REPO, the vendor package (model.safetensors + ae.safetensors + +config.json), is the default and the download source. Its files are +byte-identical to the bf16 DiT and VAE of COMFY_REPO, BFL's single-file +ComfyUI repack, and are stored under the repack's names; a repack file already +in the local ComfyUI models folder (int8_convrot = the toolkit's convrot8 +storage, attached as-is; bf16; nvfp4) is used instead of the package's. +Nothing is downloaded from the repack unless it is named as `name_or_path`. +The repack's ``_sd_`` files are the step-distilled flow-map model: never +candidates, and a loaded one is rejected by its ``ref_time_in`` keys. +""" + +BASE_REPO = "black-forest-labs/flux-3-image-klein-base" +COMFY_REPO = "black-forest-labs/flux-3-image-klein-comfy" + +# repo-relative comfy DiT files; the resolver orders them by the requested +# qtype (convrot8 -> int8, nvfp4 -> nvfp4, anything else -> bf16). The text +# encoder stays on the Comfy-Org Qwen3-VL repack (src/text_encoder.py). +COMFY_TRANSFORMER_FILES = [ + "diffusion_models/flux3_image_klein_int8_convrot.safetensors", + "diffusion_models/flux3_image_klein_bf16.safetensors", + "diffusion_models/flux3_image_klein_nvfp4.safetensors", +] +# same weights as Comfy-Org's FLUX.2 VAE file of the same name (that one is in +# the diffusers key layout; the loader accepts both) +COMFY_VAE_FILE = "vae/flux2-vae.safetensors" + +# vendor package file -> its comfy-layout location (same bytes as the repack) +PACKAGE_TO_COMFY = { + "model.safetensors": "diffusion_models/flux3_image_klein_bf16.safetensors", + "ae.safetensors": COMFY_VAE_FILE, +} diff --git a/extensions_built_in/diffusion_models/flux3/src/transformer.py b/extensions_built_in/diffusion_models/flux3/src/transformer.py index 54db610..f9e6fa9 100644 --- a/extensions_built_in/diffusion_models/flux3/src/transformer.py +++ b/extensions_built_in/diffusion_models/flux3/src/transformer.py @@ -34,6 +34,8 @@ from torch.utils.checkpoint import checkpoint from toolkit.models.v2._mixin import OstrisModelMixin +from .checkpoints import COMFY_REPO, COMFY_TRANSFORMER_FILES + try: from flash_attn_interface import flash_attn_func @@ -383,7 +385,40 @@ class SingleStreamBlock(nn.Module): return ctx + mod_gate_ctx * ctx_out, x + mod_gate_x * x_out +_QKV = ("q_proj", "k_proj", "v_proj") + + class Flux3Transformer(nn.Module, OstrisModelMixin): + aitk_comfy_repo = COMFY_REPO + aitk_comfy_weight_names = {COMFY_REPO: COMFY_TRANSFORMER_FILES} + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + if any(k.startswith("ref_time_in.") for k in state_dict): + raise ValueError( + "this is the step-distilled (flow-map) FLUX.3 checkpoint; use the " + "base model (flux3_image_klein_* without _sd_)" + ) + from toolkit.util.comfy_quant_import import split_fused_quantized_keys + + state_dict = dict(state_dict) + # the int8 repack fuses every block's q_proj/k_proj/v_proj into one qkv + # entry (rows q, k, v); the module (and the bf16/nvfp4 files) keep them + # split + for marker in [k for k in list(state_dict) if k.endswith(".qkv.comfy_quant")]: + base = marker[: -len(".qkv.comfy_quant")] + split_fused_quantized_keys( + state_dict, f"{base}.qkv", [f"{base}.{p}" for p in _QKV] + ) + for key in [k for k in list(state_dict) if k.endswith(".qkv.weight")]: + base = key[: -len(".qkv.weight")] + for p, w in zip(_QKV, torch.chunk(state_dict.pop(key), 3, dim=0)): + state_dict[f"{base}.{p}.weight"] = w.contiguous() + # static activation scales (calibration extras) are unused + for k in [k for k in state_dict if k.endswith(".input_scale")]: + state_dict.pop(k) + return state_dict + @classmethod def get_transformer_block_names(cls): return [ diff --git a/toolkit/models/v2/_mixin.py b/toolkit/models/v2/_mixin.py index bb6da09..3446161 100644 --- a/toolkit/models/v2/_mixin.py +++ b/toolkit/models/v2/_mixin.py @@ -42,6 +42,30 @@ class BasicModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMix pass +def _markers_from_header( + file_path: str, state_dict: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """ComfyUI's header-style quantized layout keeps the per-layer configs in + the safetensors metadata (``_quantization_metadata``) instead of + ``.comfy_quant`` marker tensors; synthesize the markers so the + importer sees one layout.""" + import json + + from safetensors import safe_open + + with safe_open(file_path, "pt") as f: + raw = (f.metadata() or {}).get("_quantization_metadata") + if raw is None: + return state_dict + for name, conf in json.loads(raw).get("layers", {}).items(): + marker = f"{name}.comfy_quant" + if marker not in state_dict and f"{name}.weight" in state_dict: + state_dict[marker] = torch.tensor( + list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8 + ) + return state_dict + + def adopt_component(module: torch.nn.Module, wrapper_cls: type) -> torch.nn.Module: """Rebind an already-loaded component instance onto its v2 wrapper class (in-place class swap, like the OstrisLinear conversion). Used when a @@ -584,7 +608,7 @@ class OstrisModelMixin: **kwargs, ): cls._readahead(file_path) - state_dict = load_file(file_path) + state_dict = _markers_from_header(file_path, load_file(file_path)) return cls.load_from_state_dict( state_dict, dtype,