Load comfy files if they exist in flux.3 image

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/<category>/<local name>`` (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/<rel>``
+        (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 <local_dir>/.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
 

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,
+}

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 [

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
+    ``<layer>.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,