Merge pull request #1078 from ostris/flux3-image

Add Support for FLUX.3 Image Klein Base

extensions_built_in/diffusion_models/__init__.py

@@ -23,6 +23,7 @@ from .iris3b import Iris3BModel
 from .boogu_image import BooguImageModel, BooguImageEditModel
 from .mageflow import MageFlowModel, MageFlowEditModel
 from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel
+from .flux3 import Flux3Model
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
@@ -65,4 +66,5 @@ AI_TOOLKIT_MODELS = [
     MinimaxH3Model,
     MinimaxH3Ref2VAModel,
     MinimaxH3FastModel,
+    Flux3Model,
 ]

extensions_built_in/diffusion_models/flux3/__init__.py

@@ -0,0 +1,3 @@
+from .flux3 import Flux3Model
+
+__all__ = ["Flux3Model"]

extensions_built_in/diffusion_models/flux3/flux3.py

@@ -0,0 +1,631 @@
+"""FLUX.3 (small, ~5.5B) for ai-toolkit.
+
+A joint flow-matching DiT (``src/transformer.py``) that trains and samples
+t2i and i2i (reference images as clean conditioning tokens).
+
+  - image VAE: FLUX.2 packed AE, f16, 128 latent channels
+  - text: Qwen3-VL-4B, 8 hidden layers stacked along the channel dim ->
+    ``(B, L, 20480)``; the DiT's ``timesteps_ctx`` are zeros and there is no
+    pooled vector. Embeds are stored per-sample as 2D ``(L, 20480)`` tensors
+    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):
+
+    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.
+
+Flow convention matches ai-toolkit exactly (t=1 noise -> t=0 clean,
+target = noise - clean), so no time flips anywhere.
+"""
+
+import json
+import math
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from safetensors.torch import load_file, save_file
+from torchvision.transforms.functional import to_tensor
+
+from toolkit.accelerator import unwrap_model
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+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.paths import MODELS_PATH
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.image_vae import Flux3ImageVAE
+from .src.pipeline import (
+    DEFAULT_ALPHA,
+    Flux3Pipeline,
+    build_text_conditioning,
+    pack_i2i_image_condition,
+    pack_image_latents,
+)
+from .src.text_encoder import (
+    DEFAULT_MAX_TOKENS,
+    DEFAULT_PAD_MULTIPLE,
+    QWEN3_VL_4B_REPO,
+    Flux3TextEncoder,
+    encode_flux3_prompt,
+)
+from .src.transformer import Flux3Transformer, Flux3TransformerParams
+
+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",
+}
+
+# Static flow-matching timeshift used for the training timestep/noise
+# schedule (the package's sampler alpha).
+DEFAULT_TIMESTEP_SHIFT = DEFAULT_ALPHA
+
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "use_dynamic_shifting": False,
+}
+
+
+class Flux3Model(BaseModel):
+    arch = "flux3"
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        self.target_lora_modules = ["Flux3Transformer"]
+        self.supports_model_paths = True
+        self.use_old_lokr_format = False
+
+        self.processor = None  # Qwen3-VL AutoProcessor
+        self.max_text_length = int(
+            self.model_config.model_kwargs.get("max_text_length", DEFAULT_MAX_TOKENS)
+        )
+        self.text_pad_multiple = int(
+            self.model_config.model_kwargs.get("text_pad_multiple", DEFAULT_PAD_MULTIPLE)
+        )
+        self.package_config = None
+        self._package_config_checked = False
+        # references are area-matched to the target by default; off = keep their
+        # own size, only scaling down when larger than the target area
+        self.match_target_res = bool(
+            self.model_config.model_kwargs.get("match_target_res", True)
+        )
+        self.latent_space_version = "flux3_v1"
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(
+            shift=DEFAULT_TIMESTEP_SHIFT, **scheduler_config
+        )
+
+    def get_bucket_divisibility(self):
+        return 16  # f8 conv + 2x2 packing
+
+    # ------------------------------------------------------------------
+    # Package files
+    # ------------------------------------------------------------------
+    def _package_file(self, filename: str, required: bool = True) -> Optional[str]:
+        """A package file: model_kwargs ``<stem>_path`` override, else the file
+        inside a local package dir, else its ComfyUI-layout spot under
+        MODELS_PATH (weights), else downloaded from the hub repo."""
+        override = self.model_config.model_kwargs.get(
+            filename.split(".")[0] + "_path", None
+        )
+        if override is not None:
+            if not os.path.exists(override):
+                raise FileNotFoundError(f"FLUX.3 {filename} override not found: {override}")
+            return override
+        base = self.model_config.name_or_path
+        if os.path.isdir(base):
+            path = os.path.join(base, filename)
+            if os.path.exists(path):
+                return path
+            if required:
+                raise FileNotFoundError(f"FLUX.3 package file not found: {path}")
+            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:
+            return self._comfy_weight(filename, base)
+        from huggingface_hub import hf_hub_download
+
+        try:
+            return hf_hub_download(
+                repo_id=base, filename=filename, token=os.getenv("HF_TOKEN", None)
+            )
+        except Exception as e:
+            if required:
+                raise
+            self.print_and_status_update(f" - no {filename} in {base} ({type(e).__name__})")
+            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),
+        )
+        if found is not None:
+            return found
+        from huggingface_hub import hf_hub_download
+
+        self.print_and_status_update(
+            f"Downloading {filename} from {repo_id} to {target}"
+        )
+        os.makedirs(os.path.dirname(target), exist_ok=True)
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            token=os.getenv("HF_TOKEN", None),
+            local_dir=os.path.dirname(target),
+        )
+        os.replace(path, target)
+        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)
+            if path is not None:
+                with open(path, "r") as f:
+                    self.package_config = json.load(f)
+            else:
+                self.print_and_status_update(" - using the canonical image-klein DiT layout")
+        return self.package_config
+
+    def _transformer_params(self) -> Flux3TransformerParams:
+        cfg = self._load_package_config()
+        if cfg is None:
+            return Flux3TransformerParams()
+        d = cfg["dit"]
+        in_channels = {m: c for m, c in d["in_channels"].items() if m != "text"}
+        return Flux3TransformerParams(
+            in_channels=in_channels,
+            sequence={f"x_{m}": m for m in in_channels},
+            vec_in_dim=d.get("vec_in_dim", None),
+            context_in_dim=d["in_channels"]["text"],
+            hidden_size=d["hidden"],
+            num_heads=d["heads"],
+            depth=d["depth"],
+            depth_single_blocks=d["depth_single"],
+            axes_dim=list(d["axes_dim"]),
+            theta=d["theta"],
+            mlp_ratio=d["mlp_ratio"],
+            out_modes=list(d["out_modes"]),
+        )
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    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):
+            raise ValueError(
+                "FLUX.3 package variants with pre-quantized weights "
+                f"({cfg['quantization']}) are not supported; use the bf16 package "
+                "and let the toolkit quantize."
+            )
+        if cfg is not None and cfg["dit"].get("flow_map", False):
+            raise ValueError(
+                "FLUX.3 step-distilled (flow-map) packages are not supported; "
+                "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]})"
+            )
+        del state_dict
+        flush()
+
+        attn_backend = self.model_config.model_kwargs.get("attention_backend", None)
+        if attn_backend is not None:
+            transformer.set_attention_backend(attn_backend)
+        return transformer
+
+    def _load_text_encoder(self):
+        te_path = self.model_config.te_name_or_path
+        if te_path is None:
+            te_path = self.model_config.model_kwargs.get(
+                "text_encoder_path", QWEN3_VL_4B_REPO
+            )
+        self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
+        te_kwargs = self.component_load_kwargs("te")
+        text_encoder = Flux3TextEncoder.load_model(
+            te_path,
+            dtype=te_kwargs["dtype"],
+            subfolder="",
+            qtype=te_kwargs["qtype"],  # comfy candidate ranking only
+            quantize_on_load=False,
+        )
+        processor_source = te_path if os.path.isdir(te_path) else None
+        processor = Flux3TextEncoder.load_processor(processor_source)
+        processor.tokenizer.padding_side = "right"
+        # text-only conditioning -- drop the vision tower
+        text_encoder.drop_vision_tower()
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        flush()
+        return processor, text_encoder
+
+    def _load_vae(self) -> Flux3ImageVAE:
+        self.print_and_status_update("Loading image VAE")
+        vae = Flux3ImageVAE()
+        vae.load_state_dict(load_file(self._package_file("ae.safetensors")))
+        vae.to(self.vae_torch_dtype).eval().requires_grad_(False)
+        return vae
+
+    def load_model(self):
+        self.print_and_status_update("Loading FLUX.3 model")
+
+        transformer = self._load_transformer()
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
+        flush()
+
+        processor, text_encoder = self._load_text_encoder()
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        flush()
+
+        vae = self._load_vae()
+        vae.to(self.vae_device_torch)
+
+        self.noise_scheduler = Flux3Model.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = processor.tokenizer
+        self.processor = processor
+        self.model = transformer
+        self.pipeline = Flux3Pipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Generation (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return Flux3Pipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: Flux3Pipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: AdvancedPromptEmbeds,
+        unconditional_embeds: AdvancedPromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        sc = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // sc * sc)
+        gen_config.height = int(gen_config.height // sc * sc)
+
+        # cond and uncond were encoded separately, so their pad-to-80 windows
+        # can differ. The DiT attends to real pad-token states, not zeros, so
+        # re-encode the shorter prompt at the longer window when the TE is up.
+        if unconditional_embeds is not None and self.text_encoder is not None:
+            cond_len = conditional_embeds.text_embeds[0].shape[0]
+            uncond_len = unconditional_embeds.text_embeds[0].shape[0]
+            if cond_len != uncond_len:
+                target = max(cond_len, uncond_len)
+                if cond_len < target:
+                    conditional_embeds = self._encode_at_length(gen_config.prompt, target)
+                else:
+                    unconditional_embeds = self._encode_at_length(
+                        gen_config.negative_prompt or "", target
+                    )
+
+        # references: ctrl_img_1 mirrors ctrl_img, so dedupe by path order
+        ctrl_paths = []
+        for path in (
+            gen_config.ctrl_img,
+            gen_config.ctrl_img_1,
+            gen_config.ctrl_img_2,
+            gen_config.ctrl_img_3,
+        ):
+            if path is not None and path not in ctrl_paths:
+                ctrl_paths.append(path)
+        target_pixels = gen_config.width * gen_config.height
+        ctrl_imgs = [
+            self._fit_reference(
+                to_tensor(Image.open(path).convert("RGB")).unsqueeze(0), target_pixels
+            )
+            for path in ctrl_paths
+        ]
+
+        alpha = self.model_config.model_kwargs.get("sample_alpha", None)
+        alpha = float(alpha) if alpha is not None else None
+
+        result = pipeline(
+            conditional_embeds=conditional_embeds,
+            unconditional_embeds=unconditional_embeds,
+            height=gen_config.height,
+            width=gen_config.width,
+            num_inference_steps=gen_config.num_inference_steps,
+            guidance_scale=gen_config.guidance_scale,
+            latents=gen_config.latents,
+            generator=generator,
+            ctrl_imgs=ctrl_imgs,
+            alpha=alpha,
+        )
+        return result[0]
+
+    def _fit_reference(self, img: torch.Tensor, target_pixels: int) -> torch.Tensor:
+        """Size a (1, C, H, W) reference for the DiT. References keep their own
+        aspect ratio on their own position grid; match_target_res scales the
+        area to the target's, otherwise the image only scales down when it is
+        larger than the target area. Snapped to the 16-px grid."""
+        sc = self.get_bucket_divisibility()
+        h, w = img.shape[-2], img.shape[-1]
+        if self.match_target_res or h * w > target_pixels:
+            ratio = h / w
+            new_h = math.sqrt(target_pixels * ratio)
+            new_w = new_h / ratio
+        else:
+            new_h, new_w = float(h), float(w)
+        new_h = max(sc, int(round(new_h / sc)) * sc)
+        new_w = max(sc, int(round(new_w / sc)) * sc)
+        if (new_h, new_w) != (h, w):
+            img = F.interpolate(img.float(), size=(new_h, new_w), mode="bilinear")
+        return img
+
+    def _encode_at_length(self, prompt: str, length: int) -> AdvancedPromptEmbeds:
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+        ctx = encode_flux3_prompt(
+            self.text_encoder,
+            self.processor,
+            [prompt],
+            max_length=length,
+            dtype=self.torch_dtype,
+            pad_multiple=None,
+        )
+        return AdvancedPromptEmbeds(text_embeds=[ctx[0]])
+
+    # ------------------------------------------------------------------
+    # Training hooks
+    # ------------------------------------------------------------------
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 128, h, w)
+        timestep: torch.Tensor,  # 0..1000 scale
+        text_embeddings: AdvancedPromptEmbeds,
+        batch: "DataLoaderBatchDTO" = None,
+        **kwargs,
+    ):
+        device = self.device_torch
+        dtype = self.torch_dtype
+        if self.model.device == torch.device("cpu"):
+            self.model.to(device)
+
+        batch_size = latent_model_input.shape[0]
+
+        with torch.no_grad():
+            # toolkit timestep (0..1000, 1000 = pure noise) -> flow time in [0, 1]
+            t = timestep.to(device, dtype=torch.float32) / 1000.0
+            if t.dim() == 0:
+                t = t.unsqueeze(0)
+            if t.shape[0] != batch_size:
+                t = t.expand(batch_size)
+
+            ctx, ctx_ids, ts_ctx = build_text_conditioning(
+                text_embeddings.text_embeds, device, dtype
+            )
+
+            latent_model_input = latent_model_input.to(device, dtype)
+            _, _, h_lat, w_lat = latent_model_input.shape
+            tokens, ids = pack_image_latents(latent_model_input)
+            model_kwargs = {
+                "x_image": tokens.to(dtype),
+                "x_image_ids": ids,
+                "x_image_timesteps": t.to(dtype)[:, None]
+                * torch.ones(tokens.shape[:2], device=device, dtype=dtype),
+            }
+
+            # i2i: reference image(s) as clean t=0 tokens at l = -1, -2, ...
+            control_list = None
+            if batch is not None:
+                if batch.control_tensor_list is not None:
+                    control_list = batch.control_tensor_list
+                elif batch.control_tensor is not None:
+                    control_list = [[batch.control_tensor[b]] for b in range(batch_size)]
+            if control_list is not None:
+                target_pixels = h_lat * w_lat * 16 * 16
+                ref_latents = []
+                for controls in control_list:
+                    sample_refs = []
+                    for ctrl in controls:
+                        if ctrl.ndim == 3:
+                            ctrl = ctrl.unsqueeze(0)
+                        ctrl = self._fit_reference(ctrl, target_pixels)
+                        # control tensors arrive in [0, 1]
+                        lat = self.encode_images(
+                            ctrl * 2 - 1, device=device, dtype=dtype
+                        )
+                        sample_refs.append(lat[0])
+                    ref_latents.append(sample_refs)
+                xc, xc_ids, xc_ts = pack_i2i_image_condition(ref_latents, device, dtype)
+                model_kwargs["x_image_cond"] = xc
+                model_kwargs["x_image_cond_ids"] = xc_ids
+                model_kwargs["x_image_cond_timesteps"] = xc_ts
+
+        pred = self.transformer(
+            ctx=ctx,
+            ctx_ids=ctx_ids,
+            timesteps_ctx=ts_ctx,
+            **model_kwargs,
+        )
+        out = pred["x_image"]
+        return out.reshape(batch_size, h_lat, w_lat, -1).permute(0, 3, 1, 2)
+
+    def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+
+        ctx = encode_flux3_prompt(
+            self.text_encoder,
+            self.processor,
+            prompt,
+            max_length=self.max_text_length,
+            dtype=self.torch_dtype,
+            pad_multiple=self.text_pad_multiple,
+        )
+        # one 2D (L, 20480) tensor per batch item
+        return AdvancedPromptEmbeds(text_embeds=[ctx[i] for i in range(ctx.shape[0])])
+
+    def get_loss_target(self, *args, **kwargs):
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # VAE encode / decode
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def encode_images(self, image_list, device=None, dtype=None):
+        """Images (C, H, W) in [-1, 1] -> latents (B, 128, H/16, W/16).
+
+        Accepts a list of per-item tensors or a stacked batch tensor.
+        """
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+
+        vae_dtype = self.vae.dtype
+        batch = torch.stack([img.to(device, dtype=vae_dtype) for img in image_list])
+        latents = self.vae.encode(batch)
+        return latents.to(device, dtype=dtype)
+
+    @torch.no_grad()
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        images = self.vae.decode(latents.to(device, self.vae.dtype))
+        return images.to(device, dtype)
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        from toolkit.util.quantize import dequantize_if_quantized
+
+        if not output_path.endswith(".safetensors"):
+            output_path = output_path + ".safetensors"
+        transformer: Flux3Transformer = unwrap_model(self.model)
+        state_dict = transformer.state_dict()
+        save_dict = {}
+        for k, v in state_dict.items():
+            save_dict[k] = (
+                dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype)
+            )
+        meta = get_meta_for_safetensors(meta, name=self.arch)
+        save_file(save_dict, output_path, metadata=meta)
+
+    def get_base_model_version(self):
+        return "flux3"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return Flux3Transformer.get_transformer_block_names()
+
+    def get_quantization_exclude_modules(self):
+        # sensitive small modules kept in full precision (fnmatch patterns):
+        #   emb_in*/txt_in       - token input projections
+        #   time_in*             - timestep embedder; feeds every modulation
+        #   *modulations*        - per-stream adaLN modulations
+        #   final_layer*         - modulated output projections
+        return [
+            "emb_in*",
+            "txt_in",
+            "time_in*",
+            "*modulations*",
+            "final_layer*",
+        ]
+
+    def convert_lora_weights_before_save(self, state_dict):
+        return {
+            k.replace("transformer.", "diffusion_model."): v
+            for k, v in state_dict.items()
+        }
+
+    def convert_lora_weights_before_load(self, state_dict):
+        return {
+            k.replace("diffusion_model.", "transformer."): v
+            for k, v in state_dict.items()
+        }

extensions_built_in/diffusion_models/flux3/src/image_vae.py

@@ -0,0 +1,361 @@
+"""FLUX.3 image VAE -- the FLUX.2 packed image autoencoder.
+
+f8 conv autoencoder with 32 z-channels whose mean latent is 2x2
+space-to-channel packed to ``(B, 128, H/16, W/16)`` and normalized by a
+frozen BatchNorm's running stats. ``encode``/``decode`` operate on
+``(B, 3, H, W)`` pixels in ``[-1, 1]``.
+"""
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+
+@dataclass
+class Flux3ImageVAEParams:
+    resolution: int = 256
+    in_channels: int = 3
+    ch: int = 128
+    out_ch: int = 3
+    ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4])
+    num_res_blocks: int = 2
+    z_channels: int = 32
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(
+            num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+        )
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.norm2 = nn.GroupNorm(
+            num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+        )
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels, in_channels, kernel_size=3, stride=2, padding=0
+        )
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels, in_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.conv_in = nn.Conv2d(
+            in_channels, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        self.norm_out = nn.GroupNorm(
+            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+        )
+        self.conv_out = nn.Conv2d(
+            block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        h = self.quant_conv(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        self.conv_in = nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)
+
+        self.norm_out = nn.GroupNorm(
+            num_groups=32, num_channels=block_in, eps=1e-6, affine=True
+        )
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)
+
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        h = h.to(upscale_dtype)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Flux3ImageVAE(nn.Module):
+    def __init__(self, params: Flux3ImageVAEParams | None = None):
+        super().__init__()
+        if params is None:
+            params = Flux3ImageVAEParams()
+        self.params = params
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+
+        self.bn_eps = 1e-4
+        self.bn_momentum = 0.1
+        self.ps = [2, 2]
+        self.bn = torch.nn.BatchNorm2d(
+            math.prod(self.ps) * params.z_channels,
+            eps=self.bn_eps,
+            momentum=self.bn_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def normalize(self, z: Tensor) -> Tensor:
+        self.bn.eval()
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    def encode(self, x: Tensor) -> Tensor:
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+
+        z = rearrange(
+            mean,
+            "... c (i pi) (j pj)  -> ... (c pi pj) i j",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        z = self.normalize(z)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = self.inv_normalize(z)
+        z = rearrange(
+            z,
+            "... (c pi pj) i j -> ... c (i pi) (j pj)",
+            pi=self.ps[0],
+            pj=self.ps[1],
+        )
+        dec = self.decoder(z)
+        return dec

extensions_built_in/diffusion_models/flux3/src/pipeline.py

@@ -0,0 +1,225 @@
+"""FLUX.3 sampling pipeline + shared latent/conditioning packers.
+
+Covers t2i and i2i (reference images). All conditioning arrives as prompt
+*embeds* (never raw text) plus optional control-image tensors.
+
+Flow convention: time runs t=1 (pure noise) -> t=0 (clean); the DiT
+predicts the velocity ``noise - clean`` -- identical to ai-toolkit's
+convention. The inference schedule is a timeshifted uniform grid
+(``timeshift(alpha, t)``); the package default alpha is 6.93.
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+from .positional import batched_prc_img, prc_img, prc_txt
+
+IMAGE_ROPE_SCALE = 0.5
+I2I_COND_L = -1
+DEFAULT_ALPHA = 6.93
+DEFAULT_GUIDANCE = 3.5
+
+
+def timeshift(alpha: float, t: torch.Tensor) -> torch.Tensor:
+    return alpha * t / (1.0 + (alpha - 1.0) * t)
+
+
+def build_timesteps(n_steps: int, alpha: float) -> torch.Tensor:
+    return timeshift(alpha, torch.linspace(1.0, 0.0, n_steps + 1))
+
+
+# ---------------------------------------------------------------------------
+# Text conditioning
+# ---------------------------------------------------------------------------
+
+
+def build_text_conditioning(
+    text_embeds_list: List[torch.Tensor],
+    device: torch.device,
+    dtype: torch.dtype,
+):
+    """Per-sample (L, 20480) ctx tensors -> batched DiT text kwargs.
+
+    Embeds are padded to a multiple of 80 tokens at encode time, so items
+    encoded separately (TE cache, cond/uncond pairs) can differ in length;
+    shorter ones are zero-padded to the batch's longest. Returns
+    ``(ctx (B, L, D), ctx_ids (B, L, 4), timesteps_ctx (B, L))``.
+    """
+    max_len = max(e.shape[0] for e in text_embeds_list)
+    ctx = torch.stack(
+        [
+            F.pad(e.to(device, dtype), (0, 0, 0, max_len - e.shape[0]))
+            for e in text_embeds_list
+        ]
+    )
+    _, ctx_ids = prc_txt(ctx[0])
+    ctx_ids = ctx_ids.unsqueeze(0).expand(ctx.shape[0], -1, -1).to(device)
+    timesteps_ctx = torch.zeros(ctx.shape[:2], device=device, dtype=dtype)
+    return ctx, ctx_ids, timesteps_ctx
+
+
+# ---------------------------------------------------------------------------
+# Token packing (training + sampling share these)
+# ---------------------------------------------------------------------------
+
+
+def pack_image_latents(latents: torch.Tensor):
+    """(B, 128, h, w) image latents -> ((B, h*w, 128) tokens, float ids with
+    the image RoPE scale applied to the h/w axes)."""
+    tokens, ids = batched_prc_img(latents)
+    ids = ids.to(dtype=torch.float32)
+    ids[..., 1:3] = ids[..., 1:3] * IMAGE_ROPE_SCALE
+    return tokens, ids
+
+
+def pack_i2i_image_condition(
+    latents_per_sample: List[List[torch.Tensor]],
+    device: torch.device,
+    dtype: torch.dtype,
+):
+    """Reference-image latents for i2i editing.
+
+    ``latents_per_sample`` is one list of (128, h, w) latents per batch item
+    (every item must have the same reference geometry). Reference i placed at
+    ``t = 0, l = -1 - i``; conditioning tokens run at timestep 0.
+    Returns ``(tokens, ids, timesteps)``.
+    """
+    all_tokens, all_ids = [], []
+    for refs in latents_per_sample:
+        parts, id_parts = [], []
+        for i_ref, latent in enumerate(refs):
+            x_ref, x_ref_ids = prc_img(
+                latent,
+                t_coord=torch.tensor([0], device=latent.device, dtype=torch.long),
+                l_coord=torch.tensor(
+                    [I2I_COND_L - i_ref], device=latent.device, dtype=torch.long
+                ),
+            )
+            parts.append(x_ref)
+            id_parts.append(x_ref_ids)
+        all_tokens.append(torch.cat(parts, dim=0))
+        all_ids.append(torch.cat(id_parts, dim=0))
+    tokens = torch.stack(all_tokens).to(device, dtype)
+    ids = torch.stack(all_ids).to(device=device, dtype=torch.float32)
+    ids[..., 1:3] = ids[..., 1:3] * IMAGE_ROPE_SCALE
+    timesteps = torch.zeros(tokens.shape[:2], device=device, dtype=dtype)
+    return tokens, ids, timesteps
+
+
+# ---------------------------------------------------------------------------
+# Sampling pipeline
+# ---------------------------------------------------------------------------
+
+
+class Flux3Pipeline:
+    """Lightweight sampler used by ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # ``model`` is the Flux3Model so we can reuse its encode/decode + config.
+        self.model = model
+
+    @property
+    def device(self):
+        return self.model.device_torch
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def set_progress_bar_config(self, **kwargs):
+        pass
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 50,
+        guidance_scale: float = DEFAULT_GUIDANCE,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        ctrl_imgs: Optional[List[torch.Tensor]] = None,  # each (1, 3, H, W) in [0, 1]
+        alpha: Optional[float] = None,
+        **kwargs,
+    ):
+        model = self.model
+        device = model.device_torch
+        dtype = model.torch_dtype
+        transformer = model.transformer
+
+        do_cfg = unconditional_embeds is not None and guidance_scale > 1.0
+
+        if alpha is None:
+            alpha = DEFAULT_ALPHA
+        timesteps = build_timesteps(num_inference_steps, alpha=alpha).to(device)
+
+        # text conditioning, CFG-doubled as [uncond, cond]
+        cond_list = conditional_embeds.text_embeds
+        if do_cfg:
+            embeds_list = unconditional_embeds.text_embeds + cond_list
+        else:
+            embeds_list = cond_list
+        ctx, ctx_ids, ts_ctx = build_text_conditioning(embeds_list, device, dtype)
+        cfg_batch = ctx.shape[0]
+
+        # initial noise tokens
+        h_lat, w_lat = height // 16, width // 16
+        if latents is None:
+            latents = randn_tensor(
+                (1, 128, h_lat, w_lat), generator=generator, dtype=torch.float32
+            )
+        latents = latents.to(device)
+        x, ids = pack_image_latents(latents)
+        x = x.float()
+        ids = ids.repeat(cfg_batch, 1, 1).to(device)
+
+        # reference-image conditioning: reference i rides at l = -1 - i
+        cond_kwargs = {}
+        if ctrl_imgs:
+            ref_latents = []
+            for ctrl_img in ctrl_imgs:
+                ctrl = ctrl_img.to(device, dtype) * 2 - 1
+                ref_latents.append(
+                    model.encode_images(list(ctrl), device=device, dtype=dtype)[0]
+                )
+            xc, xc_ids, xc_ts = pack_i2i_image_condition([ref_latents], device, dtype)
+            cond_kwargs = {
+                "x_image_cond": xc.repeat(cfg_batch, 1, 1),
+                "x_image_cond_ids": xc_ids.repeat(cfg_batch, 1, 1),
+                "x_image_cond_timesteps": xc_ts.repeat(cfg_batch, 1),
+            }
+
+        # Euler integration of the flow ODE with CFG
+        for step_index, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
+            x_in = x.to(dtype).repeat(cfg_batch, 1, 1)
+            pred = transformer(
+                ctx=ctx,
+                ctx_ids=ctx_ids,
+                timesteps_ctx=ts_ctx,
+                x_image=x_in,
+                x_image_ids=ids,
+                x_image_timesteps=t_curr
+                * torch.ones(x_in.shape[:2], device=device, dtype=dtype),
+                **cond_kwargs,
+            )
+            v = pred["x_image"].float()
+            if do_cfg:
+                v_uc, v_c = v.chunk(2)
+                v = v_uc + guidance_scale * (v_c - v_uc)
+            if getattr(model, "sample_step_hook", None) is not None:
+                # inference engine preview: x0 estimate
+                x0 = (x - t_curr * v).reshape(1, h_lat, w_lat, -1).permute(0, 3, 1, 2)
+                model._emit_sample_step(x0, step_index, len(timesteps) - 1)
+            x = x + (t_prev - t_curr) * v
+
+        img_latent = x.reshape(1, h_lat, w_lat, -1).permute(0, 3, 1, 2)
+        images = model.decode_latents(img_latent, device=device, dtype=dtype)
+        images = images.float().clamp(-1.0, 1.0)
+        images = ((images + 1.0) * 127.5).round().to(torch.uint8)
+        images = images.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(arr) for arr in images]

extensions_built_in/diffusion_models/flux3/src/positional.py

@@ -0,0 +1,69 @@
+"""Position-id packers for the FLUX.3 joint sequence.
+
+Every token stream fed to the DiT carries a 4-axis position id ``(t, h, w, l)``.
+Image tokens use ``h``/``w`` (``t = 0``); the ``l`` axis separates reference
+images (i2i conditioning uses ``l = -1, -2, ...``) and carries the sequence
+position for text tokens.
+"""
+
+from typing import Callable
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def prc_img(
+    x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None
+) -> tuple[Tensor, Tensor]:
+    """(C, H, W) latent -> ((H*W, C) tokens, (H*W, 4) ids)."""
+    _, h, w = x.shape
+    dev = x.device
+    coords = {
+        "t": torch.arange(1, device=dev) if t_coord is None else t_coord.to(dev),
+        "h": torch.arange(h, device=dev),
+        "w": torch.arange(w, device=dev),
+        "l": torch.arange(1, device=dev) if l_coord is None else l_coord.to(dev),
+    }
+    x_ids = torch.cartesian_prod(coords["t"], coords["h"], coords["w"], coords["l"])
+    x = rearrange(x, "c h w -> (h w) c")
+    return x, x_ids
+
+
+def prc_txt(
+    x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None
+) -> tuple[Tensor, Tensor]:
+    """(L, D) text context -> ((L, D) tokens, (L, 4) ids) with l = arange(L)."""
+    assert l_coord is None, "l_coord not supported for txt"
+    length, _ = x.shape
+    dev = x.device
+    coords = {
+        "t": torch.arange(1, device=dev) if t_coord is None else t_coord.to(dev),
+        "h": torch.arange(1, device=dev),
+        "w": torch.arange(1, device=dev),
+        "l": torch.arange(length, device=dev),
+    }
+    x_ids = torch.cartesian_prod(coords["t"], coords["h"], coords["w"], coords["l"])
+    return x, x_ids
+
+
+def _batched(fn: Callable) -> Callable:
+    def wrapped(
+        x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None
+    ) -> tuple[Tensor, Tensor]:
+        pairs = [
+            fn(
+                x[i],
+                t_coord[i] if t_coord is not None else None,
+                l_coord[i] if l_coord is not None else None,
+            )
+            for i in range(len(x))
+        ]
+        xs, ids = zip(*pairs)
+        return torch.stack(xs), torch.stack(ids)
+
+    return wrapped
+
+
+batched_prc_img = _batched(prc_img)
+batched_prc_txt = _batched(prc_txt)

extensions_built_in/diffusion_models/flux3/src/text_encoder.py

@@ -0,0 +1,123 @@
+"""FLUX.3 text conditioning -- Qwen3-VL-4B hidden-state stack.
+
+The DiT text context is the stack of a fixed set of Qwen3-VL hidden layers
+along the channel dim: ``ctx = (B, L, len(SELECT_LAYERS) * 2560 = 20480)``.
+Prompts go through the chat template and are right-padded to a multiple of
+``pad_multiple`` tokens (the package's ``text_encoder.pad_multiple``); the
+model was trained with the pad tokens in-context and the DiT has no text
+attention mask. FLUX.3 feeds the DiT zero ``timesteps_ctx`` and no pooled
+vector.
+"""
+
+import math
+from typing import List, Optional
+
+import torch
+from einops import rearrange
+
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder
+
+SELECT_LAYERS = [4, 8, 12, 16, 20, 24, 28, 32]
+CONTEXT_DIM = 20480  # len(SELECT_LAYERS) * qwen hidden (2560)
+
+# package defaults (text_encoder section of config.json)
+DEFAULT_PAD_MULTIPLE = 80
+DEFAULT_MAX_TOKENS = 8192
+
+QWEN3_VL_4B_REPO = "Qwen/Qwen3-VL-4B-Instruct"
+
+
+class Flux3TextEncoder(Qwen3VLTextEncoder):
+    """Qwen3-VL-4B-Instruct. The package's text_encoder/ and tokenizer/ are
+    byte-identical to the public Qwen repo, so config + tokenizer come from
+    there and the weights from the Comfy-Org single-file repack under
+    MODELS_PATH/text_encoders/."""
+
+    aitk_config_repo = QWEN3_VL_4B_REPO
+    aitk_tokenizer_repo = QWEN3_VL_4B_REPO
+    aitk_tokenizer_subfolder = ""
+    aitk_processor_repo = QWEN3_VL_4B_REPO
+    aitk_processor_subfolder = ""
+
+    aitk_comfy_repo = "Comfy-Org/Qwen3-VL"
+    _COMFY_FILES = [
+        "text_encoders/qwen3vl_4b_bf16.safetensors",
+        "text_encoders/qwen3vl_4b_fp8_scaled.safetensors",
+    ]
+    aitk_comfy_weight_names = {QWEN3_VL_4B_REPO: _COMFY_FILES}
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        # comfy files drop the tied lm_head; strict load wants the key
+        embed = state_dict.get("model.language_model.embed_tokens.weight")
+        if embed is not None and "lm_head.weight" not in state_dict:
+            state_dict["lm_head.weight"] = embed
+        return state_dict
+
+
+def padded_token_length(
+    tokenizer,
+    texts: List[str],
+    pad_multiple: int,
+    max_tokens: int,
+) -> int:
+    """Common padded length for a batch: the longest prompt rounded up to a
+    multiple of ``pad_multiple``, capped at ``max_tokens``."""
+    longest = 0
+    for text in texts:
+        real = tokenizer(
+            text, return_tensors="pt", truncation=True, max_length=max_tokens
+        )["input_ids"].shape[1]
+        longest = max(longest, real)
+    return min(math.ceil(longest / pad_multiple) * pad_multiple, max_tokens)
+
+
+@torch.no_grad()
+def encode_flux3_prompt(
+    text_encoder,
+    processor,
+    prompts: List[str],
+    max_length: int = DEFAULT_MAX_TOKENS,
+    dtype: torch.dtype = torch.bfloat16,
+    pad_multiple: Optional[int] = DEFAULT_PAD_MULTIPLE,
+) -> torch.Tensor:
+    """Prompts -> text context ``(B, L, 20480)``.
+
+    ``text_encoder`` is a ``Qwen3VLForConditionalGeneration``; ``processor`` the
+    matching ``AutoProcessor``. Prompts are wrapped in the chat template with a
+    generation prompt and right-padded (matching how the DiT was trained -- do
+    not strip padding). ``L`` is the batch's longest prompt rounded up to a
+    multiple of ``pad_multiple``, capped at ``max_length``; with
+    ``pad_multiple=None`` every prompt is padded to exactly ``max_length``.
+    """
+    messages = [[{"role": "user", "content": p}] for p in prompts]
+    texts = [
+        processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+        for msg in messages
+    ]
+    tokenizer = processor.tokenizer
+    if pad_multiple is not None:
+        max_length = padded_token_length(tokenizer, texts, pad_multiple, max_length)
+    input_ids_list, attn_list = [], []
+    for text in texts:
+        toks = tokenizer(
+            text,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=max_length,
+        )
+        input_ids_list.append(toks["input_ids"])
+        attn_list.append(toks["attention_mask"])
+    input_ids = torch.cat(input_ids_list, dim=0).to(text_encoder.device)
+    attention_mask = torch.cat(attn_list, dim=0).to(text_encoder.device)
+
+    out = text_encoder(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        output_hidden_states=True,
+        use_cache=False,
+    )
+    stacked = torch.stack([out.hidden_states[k] for k in SELECT_LAYERS], dim=1)
+    result = rearrange(stacked, "b c l d -> b l (c d)")
+    return result.to(dtype)

extensions_built_in/diffusion_models/flux3/src/transformer.py

@@ -0,0 +1,708 @@
+"""Flux3Transformer -- the FLUX.3 joint single-sequence DiT.
+
+Every content mode is a token stream: ``x_image`` is the denoised stream and
+``x_image_cond`` carries clean (t=0) reference tokens. Each stream is a
+``(B, S, C_mode)`` token sequence with 4-axis RoPE position ids ``(B, S, 4)``
+and per-token timesteps ``(B, S)``. Text context ``ctx`` is ``(B, L, 20480)``
+(stacked Qwen3-VL hidden layers); FLUX.3 has no pooled text vector.
+
+Each mode gets its own input projection, ``depth`` early self-attention
+blocks and modulation tables; text gets the same. The active modes and text
+are then concatenated into one sequence for ``depth_single_blocks`` joint
+blocks, split back apart, and the output modes go through their own final
+layer. A mode with no tokens runs a 2-token zero dummy through its early
+blocks (keeps the graph uniform) and never enters the joint sequence.
+
+The module tree matches the shipped ``model.safetensors`` layout so the state
+dict loads without any key renaming.
+
+Attention defaults to torch's native ``scaled_dot_product_attention`` so it
+runs on any GPU; ``set_attention_backend`` can opt into the cuDNN-pinned
+kernel (Blackwell) or flash-attn 3 when installed.
+"""
+
+import math
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor, nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.utils.checkpoint import checkpoint
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+
+try:
+    from flash_attn_interface import flash_attn_func
+
+    _FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    flash_attn_func = None
+    _FLASH_ATTN_AVAILABLE = False
+
+ATTENTION_BACKENDS = ("native", "cudnn", "flash")
+
+
+@dataclass
+class Flux3TransformerParams:
+    # content modes -> latent channels (text is separate: context_in_dim)
+    in_channels: dict = field(
+        default_factory=lambda: {"image": 128, "image_cond": 128}
+    )
+    # forward kwarg name -> content mode
+    sequence: dict = field(
+        default_factory=lambda: {"x_image": "image", "x_image_cond": "image_cond"}
+    )
+    # pooled-vector embedder width; None = no vector_in
+    vec_in_dim: int | None = None
+    context_in_dim: int = 20480
+    hidden_size: int = 3072
+    num_heads: int = 24
+    depth: int = 5
+    depth_single_blocks: int = 28
+    axes_dim: list[int] = field(default_factory=lambda: [32, 32, 32, 32])
+    theta: int = 10000
+    mlp_ratio: float = 3.0
+    qkv_bias: bool = False
+    attn_mode: str = "native"
+    # modes with a final_layer (velocity outputs); None = every content mode
+    out_modes: list[str] | None = field(default_factory=lambda: ["image"])
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack(
+        [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
+    )
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class Attention(nn.Module):
+    def forward(
+        self, q: Tensor, k: Tensor, v: Tensor, attn_mode: str = "native"
+    ) -> Tensor:
+        if attn_mode == "flash":
+            if flash_attn_func is None:
+                raise RuntimeError("flash_attn_func is unavailable")
+            b, h, n, d = q.shape
+            q, k, v = (rearrange(t, "b h n d -> b n h d") for t in (q, k, v))
+            x = flash_attn_func(q, k, v)[0]
+            return x.reshape(b, n, h * d)
+        if attn_mode == "cudnn":
+            with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+                x = F.scaled_dot_product_attention(q, k, v)
+        else:  # native
+            x = F.scaled_dot_product_attention(q, k, v)
+        return rearrange(x, "b h n d -> b n (h d)")
+
+
+def _bf16_if_required(module: nn.Module, x: Tensor) -> Tensor:
+    if getattr(module, "_requires_bf16_input", False) and x.dtype != torch.bfloat16:
+        return x.to(torch.bfloat16)
+    return x
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [
+                rope(ids[..., i], self.axes_dim[i], self.theta)
+                for i in range(len(self.axes_dim))
+            ],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(
+    t: Tensor,
+    dim: int,
+    max_period: int = 10000,
+    time_factor: float = 1000.0,
+    out_dtype: torch.dtype | None = None,
+) -> Tensor:
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(0, half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(out_dtype if out_dtype is not None else t.dtype)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, disable_bias: bool = False):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=not disable_bias)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=not disable_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class SiLUActivation(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gate_fn = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return self.gate_fn(x1) * x2
+
+
+class Modulation(nn.Module):
+    def __init__(self, dim: int, double: bool, disable_bias: bool = False):
+        super().__init__()
+        self.is_double = double
+        self.multiplier = 6 if double else 3
+        self.lin = nn.Linear(dim, self.multiplier * dim, bias=not disable_bias)
+
+    def forward(self, vec: Tensor):
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        out = out.chunk(self.multiplier, dim=-1)
+        return out[:3], out[3:] if self.is_double else None
+
+
+class LastLayer(nn.Module):
+    def __init__(self, hidden_size: int, out_channels: int):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=False)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=False)
+        )
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        if vec.ndim == 2:
+            vec = vec[:, None, :]
+        s = self.adaLN_modulation[0](vec)  # type: ignore
+        shift_w, scale_w = self.adaLN_modulation[1].weight.chunk(2)
+
+        # x = (1 + scale) * self.norm_final(x) + shift
+        # avoid materializing the full (seq, 2*hidden) projection to reduce memory
+        x = self.norm_final(x)
+        x = x * (F.linear(s, scale_w) + 1)
+        x = x + F.linear(s, shift_w)
+        return self.linear(x)
+
+
+class ModeBlock(nn.Module):
+    """Per-mode (early) block: self attention within one stream."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 3.0,
+        qkv_bias: bool = False,
+        use_gated_mlp: bool = True,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.hidden_size = hidden_size
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_mult_factor = 2 if use_gated_mlp else 1
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.mlp_in = nn.Linear(
+            hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor, bias=False
+        )
+
+        self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False)
+
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = (
+            SiLUActivation() if use_gated_mlp else nn.GELU(approximate="tanh")
+        )
+        self.attention = Attention()
+        self.attn_mode = "native"
+
+    def forward(self, x: Tensor, pe: Tensor, mod: tuple[Tensor, ...]) -> Tensor:
+        mod_shift, mod_scale, mod_gate = mod
+        x_mod = (1 + mod_scale) * self.pre_norm(x) + mod_shift
+
+        q = rearrange(
+            self.q_proj(_bf16_if_required(self.q_proj, x_mod)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        k = rearrange(
+            self.k_proj(_bf16_if_required(self.k_proj, x_mod)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        v = rearrange(
+            self.v_proj(_bf16_if_required(self.v_proj, x_mod)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        mlp = self.mlp_in(_bf16_if_required(self.mlp_in, x_mod))
+
+        q, k = self.norm(q, k, v)
+        q, k = apply_rope(q, k, pe)
+        attn = self.attention(q, k, v, self.attn_mode)
+
+        output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out(
+            _bf16_if_required(self.mlp_out, self.mlp_act(mlp))
+        )
+        return x + mod_gate * output
+
+
+class SingleStreamBlock(nn.Module):
+    """Joint block: text + all active content modes in one attention sequence."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 3.0,
+        use_gated_mlp: bool = True,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.hidden_size = hidden_size
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp_mult_factor = 2 if use_gated_mlp else 1
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.mlp_in = nn.Linear(
+            hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor, bias=False
+        )
+        self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False)
+
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = (
+            SiLUActivation() if use_gated_mlp else nn.GELU(approximate="tanh")
+        )
+        self.attention = Attention()
+        self.attn_mode = "native"
+
+    def forward(
+        self,
+        x: Tensor,
+        ctx: Tensor,
+        pe: Tensor,
+        mod_x: tuple[Tensor, ...],
+        mod_ctx: tuple[Tensor, ...],
+    ) -> tuple[Tensor, Tensor]:
+        mod_shift_x, mod_scale_x, mod_gate_x = mod_x
+        mod_shift_ctx, mod_scale_ctx, mod_gate_ctx = mod_ctx
+
+        x_mod = (1 + mod_scale_x) * self.pre_norm(x) + mod_shift_x
+        ctx_mod = (1 + mod_scale_ctx) * self.pre_norm(ctx) + mod_shift_ctx
+        combined = torch.cat((ctx_mod, x_mod), dim=1)
+
+        q = rearrange(
+            self.q_proj(_bf16_if_required(self.q_proj, combined)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        k = rearrange(
+            self.k_proj(_bf16_if_required(self.k_proj, combined)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        v = rearrange(
+            self.v_proj(_bf16_if_required(self.v_proj, combined)),
+            "B L (H D) -> B H L D",
+            H=self.num_heads,
+        )
+        mlp = self.mlp_in(_bf16_if_required(self.mlp_in, combined))
+
+        q, k = self.norm(q, k, v)
+        q, k = apply_rope(q, k, pe)
+        attn = self.attention(q, k, v, self.attn_mode)
+
+        output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out(
+            _bf16_if_required(self.mlp_out, self.mlp_act(mlp))
+        )
+
+        ctx_out, x_out = torch.split(output, [ctx.shape[1], x.shape[1]], dim=1)
+        return ctx + mod_gate_ctx * ctx_out, x + mod_gate_x * x_out
+
+
+class Flux3Transformer(nn.Module, OstrisModelMixin):
+    @classmethod
+    def get_transformer_block_names(cls):
+        return [
+            "content_mode_blocks.image",
+            "content_mode_blocks.image_cond",
+            "txt_mode_blocks",
+            "single_blocks",
+        ]
+
+    def get_offload_ignore_modules(self):
+        # modulation tables stay resident on the compute device; LastLayer
+        # reads adaLN weights raw (bypasses offload forward hijacks)
+        return [m for m in self.modules() if isinstance(m, (Modulation, LastLayer))]
+
+    def __init__(self, params: Flux3TransformerParams | None = None):
+        super().__init__()
+        if params is None:
+            params = Flux3TransformerParams()
+        self.params = params
+
+        self.sequence = params.sequence
+        for _, v in params.sequence.items():
+            assert v in params.in_channels, (params.sequence, params.in_channels)
+
+        self.in_channels = params.in_channels
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.depth = params.depth
+        self.depth_single_blocks = params.depth_single_blocks
+        self.mlp_ratio = params.mlp_ratio
+        self.qkv_bias = params.qkv_bias
+        self.attn_mode = params.attn_mode
+
+        self.gradient_checkpointing = False
+
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by "
+                f"num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(
+                f"Got {params.axes_dim} but expected positional dim {pe_dim}"
+            )
+
+        self.pe_embedder = EmbedND(
+            dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim
+        )
+
+        self.emb_in = nn.ModuleDict(
+            {
+                m: nn.Linear(c, params.hidden_size, bias=False)
+                for m, c in params.in_channels.items()
+            }
+        )
+        self.txt_in = nn.Linear(params.context_in_dim, params.hidden_size, bias=False)
+
+        self.time_in = MLPEmbedder(
+            in_dim=256, hidden_dim=params.hidden_size, disable_bias=True
+        )
+        self.vector_in: nn.Module = (
+            MLPEmbedder(params.vec_in_dim, params.hidden_size, disable_bias=True)
+            if params.vec_in_dim is not None
+            else nn.Identity()
+        )
+
+        content_modes = sorted(params.in_channels.keys())
+        self.early_stream_modulations = nn.ModuleDict(
+            {
+                **{
+                    m: Modulation(params.hidden_size, double=False, disable_bias=True)
+                    for m in content_modes
+                },
+                "txt": Modulation(params.hidden_size, double=False, disable_bias=True),
+            }
+        )
+        self.single_stream_modulations = nn.ModuleDict(
+            {
+                **{
+                    m: Modulation(params.hidden_size, double=False, disable_bias=True)
+                    for m in content_modes
+                },
+                "txt": Modulation(params.hidden_size, double=False, disable_bias=True),
+            }
+        )
+
+        self.content_mode_blocks = nn.ModuleDict(
+            {
+                m: nn.ModuleList(
+                    [
+                        ModeBlock(
+                            params.hidden_size,
+                            params.num_heads,
+                            mlp_ratio=params.mlp_ratio,
+                            qkv_bias=params.qkv_bias,
+                        )
+                        for _ in range(params.depth)
+                    ]
+                )
+                for m in content_modes
+            }
+        )
+        self.txt_mode_blocks = nn.ModuleList(
+            [
+                ModeBlock(
+                    params.hidden_size,
+                    params.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(
+                    params.hidden_size,
+                    params.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                )
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        out_modes = (
+            list(params.in_channels) if params.out_modes is None else params.out_modes
+        )
+        assert set(out_modes) <= set(params.in_channels), (out_modes, params.in_channels)
+        self.final_layer = nn.ModuleDict(
+            {
+                m: LastLayer(params.hidden_size, c)
+                for m, c in params.in_channels.items()
+                if m in out_modes
+            }
+        )
+
+        self.set_attention_backend(params.attn_mode)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def enable_gradient_checkpointing(self, enable: bool = True):
+        self.gradient_checkpointing = enable
+
+    def disable_gradient_checkpointing(self):
+        self.gradient_checkpointing = False
+
+    def set_attention_backend(self, backend: str):
+        if backend not in ATTENTION_BACKENDS:
+            raise ValueError(
+                f"Unknown attention backend {backend!r}; choose from {ATTENTION_BACKENDS}"
+            )
+        if backend == "flash" and not _FLASH_ATTN_AVAILABLE:
+            raise RuntimeError(
+                "attention backend 'flash' requested but flash_attn_interface "
+                "(flash-attn 3) is not installed"
+            )
+        if backend == "cudnn" and not torch.cuda.is_available():
+            raise RuntimeError("attention backend 'cudnn' requires CUDA")
+        self.attn_mode = backend
+        for module in self.modules():
+            if hasattr(module, "attn_mode"):
+                module.attn_mode = backend
+
+    def _ckpt(self, module, *args):
+        if self.gradient_checkpointing and torch.is_grad_enabled():
+            return checkpoint(module, *args, use_reentrant=False)
+        return module(*args)
+
+    def forward(
+        self,
+        ctx: Tensor,
+        ctx_ids: Tensor,
+        timesteps_ctx: Tensor,
+        vector: Tensor | None = None,
+        **kwargs: Tensor,
+    ) -> dict[str, Tensor]:
+        batch_size = ctx.shape[0]
+
+        x_by_mod: dict[str, list] = defaultdict(list)
+        ids_by_mod: dict[str, list] = defaultdict(list)
+        ts_by_mod: dict[str, list] = defaultdict(list)
+        lens_by_mod: dict[str, list] = defaultdict(list)
+        names_by_mod: dict[str, list] = defaultdict(list)
+        for k, v in self.sequence.items():
+            if k not in kwargs or kwargs[k] is None:
+                continue
+            assert k + "_ids" in kwargs and k + "_timesteps" in kwargs
+            x_by_mod[v].append(kwargs[k])
+            ids_by_mod[v].append(kwargs[k + "_ids"])
+            ts_by_mod[v].append(kwargs[k + "_timesteps"])
+            lens_by_mod[v].append(kwargs[k].shape[1])
+            names_by_mod[v].append(k)
+
+        active = list(x_by_mod.keys())
+        all_mods = sorted(self.in_channels.keys())
+        inactive = [m for m in all_mods if m not in active]
+        dummy_len = 2  # attention kernels need >=2
+
+        device = x_by_mod[active[0]][0].device
+        dtype = x_by_mod[active[0]][0].dtype
+        for k in all_mods:
+            if k in active:
+                x_by_mod[k] = torch.cat(x_by_mod[k], 1)
+                ids_by_mod[k] = torch.cat(ids_by_mod[k], 1)
+                ts_by_mod[k] = torch.cat(ts_by_mod[k], 1)
+            else:
+                x_by_mod[k] = torch.zeros(
+                    batch_size,
+                    dummy_len,
+                    self.in_channels[k],
+                    device=device,
+                    dtype=dtype,
+                )
+                ids_by_mod[k] = torch.zeros(
+                    batch_size, dummy_len, 4, device=device, dtype=torch.int64
+                )
+                ts_by_mod[k] = torch.zeros(
+                    batch_size, dummy_len, device=device, dtype=dtype
+                )
+
+        vec_by_mod: dict[str, Tensor] = {}
+        for k in all_mods:
+            ts = rearrange(ts_by_mod[k], "b t -> (b t)")
+            vec_by_mod[k] = rearrange(
+                self.time_in(timestep_embedding(ts, 256, out_dtype=dtype)),
+                "(b t) d -> b t d",
+                b=batch_size,
+            )
+        ts_ctx = rearrange(timesteps_ctx, "b t -> (b t)")
+        vec_ctx = rearrange(
+            self.time_in(timestep_embedding(ts_ctx, 256, out_dtype=dtype)),
+            "(b t) d -> b t d",
+            b=batch_size,
+        )
+
+        if self.params.vec_in_dim is not None:
+            assert vector is not None and vector.ndim == 2
+            vector_emb = self.vector_in(vector)
+            for k in all_mods:
+                vec_by_mod[k] = vec_by_mod[k] + vector_emb[:, None, :]
+            vec_ctx = vec_ctx + vector_emb[:, None, :]
+
+        early_content_mods: dict[str, tuple] = {}
+        for key in self.early_stream_modulations:
+            if key == "txt":
+                continue
+            mod, _ = self.early_stream_modulations[key](vec_by_mod[key])
+            early_content_mods[key] = mod
+        early_txt_mod, _ = self.early_stream_modulations["txt"](vec_ctx)
+
+        single_content_mods: dict[str, list] = {}
+        for key in self.single_stream_modulations:
+            if key == "txt":
+                continue
+            mod, _ = self.single_stream_modulations[key](vec_by_mod[key])
+            single_content_mods[key] = list(mod)
+        single_txt_mod, _ = self.single_stream_modulations["txt"](vec_ctx)
+
+        content_by_mod = {k: self.emb_in[k](x_by_mod[k]) for k in all_mods}
+        txt = self.txt_in(ctx)
+
+        pe_by_mod = {k: self.pe_embedder(ids_by_mod[k]) for k in all_mods}
+        pe_ctx = self.pe_embedder(ctx_ids)
+
+        for d in range(self.depth):
+            for m in all_mods:
+                content_by_mod[m] = self._ckpt(
+                    self.content_mode_blocks[m][d],
+                    content_by_mod[m],
+                    pe_by_mod[m],
+                    early_content_mods[m],
+                )
+            txt = self._ckpt(self.txt_mode_blocks[d], txt, pe_ctx, early_txt_mod)
+
+        img = torch.cat([content_by_mod[k] for k in active], dim=1)
+        content_mod = tuple(
+            torch.cat([single_content_mods[k][i] for k in active], dim=1)
+            for i in range(3)
+        )
+        pe_content = torch.cat([pe_by_mod[k] for k in active], dim=2)
+        pe = torch.cat((pe_ctx, pe_content), dim=2)
+
+        for block in self.single_blocks:
+            txt, img = self._ckpt(
+                block,
+                img,
+                txt,
+                pe,
+                content_mod,
+                single_txt_mod,
+            )
+
+        img_by_mod: dict[str, Tensor] = {}
+        pos = 0
+        for k in active:
+            length = sum(lens_by_mod[k])
+            img_by_mod[k] = img[:, pos : pos + length]
+            pos += length
+        for k in inactive:
+            img_by_mod[k] = content_by_mod[k]
+
+        final = {
+            k: self._ckpt(self.final_layer[k], img_by_mod[k], vec_by_mod[k])
+            for k in self.final_layer
+        }
+        out: dict[str, Tensor] = {}
+        for k in active:
+            if k not in final:
+                continue
+            pos = 0
+            for name, length in zip(names_by_mod[k], lens_by_mod[k]):
+                out[name] = final[k][:, pos : pos + length]
+                pos += length
+        return out

extensions_built_in/diffusion_models/ui.tsx

@@ -1715,6 +1715,45 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       "datasets.auto_frame_count",
     ],
   },
+  {
+    name: "flux3",
+    label: "FLUX.3 Image Klein (base)",
+    group: "image",
+    gateUrl: "https://huggingface.co/black-forest-labs/flux-3-image-klein-base",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "config.process[0].model.name_or_path": [
+        "black-forest-labs/flux-3-image-klein-base",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].sample.width": [1024, 1024],
+      "config.process[0].sample.height": [1024, 1024],
+      "config.process[0].sample.guidance_scale": [3.5, 4],
+      "config.process[0].sample.sample_steps": [28, 25],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      "config.process[0].model.model_kwargs": [
+        {
+          match_target_res: true,
+        },
+        {},
+      ],
+    },
+    disableSections: ["network.conv"],
+    additionalSections: [
+      "datasets.multi_control_paths",
+      "sample.multi_ctrl_imgs",
+      "model.layer_offloading",
+      "model.low_vram",
+      "model.qie.match_target_res",
+    ],
+  },
   {
     name: "flux2_klein_4b",
     label: "FLUX.2-klein-base-4B",

extensions_built_in/inference_engine/latent_format_tables.py

@@ -540,4 +540,45 @@ LATENT_FORMATS = {
             [0.0447, 0.0994, 0.0881],
         ],
     },
+    "Flux3Image": {
+        # FLUX.3 image VAE (FLUX.2 packed f8 autoencoder, BN-normalized): 32 ch x 2x2 shuffle.
+        # least-squares fit on flux3-small latents (2026-09-12), r2 ~ [0.74, 0.76, 0.76]
+        "channels": 128, "dims": 2, "spatial": 16, "temporal": 1,
+        "reshape": "flux2_2x2",
+        "bias": [-0.01281, -0.10207, -0.2469],
+        "factors": [
+            [0.01443, 0.01462, 0.02385],
+            [-0.01563, -0.05629, 0.03518],
+            [-0.00366, 0.01501, 0.0998],
+            [0.17186, 0.27951, 0.40425],
+            [0.05145, -0.01879, -0.02143],
+            [0.01929, -0.00069, -0.03722],
+            [-0.0942, 0.11153, -0.09854],
+            [0.03436, 0.01086, -0.00353],
+            [-0.33485, -0.22427, -0.04648],
+            [0.01356, -0.00215, 0.03131],
+            [-0.00078, 0.0166, -0.01035],
+            [0.05476, 0.04449, -0.06473],
+            [0.00259, -0.00427, -0.01085],
+            [0.02877, 0.00309, -0.02534],
+            [-0.01333, 0.0106, 0.01833],
+            [-0.0832, 0.09934, 0.07942],
+            [-0.00951, -0.00139, 0.0027],
+            [0.01677, -0.01178, -0.02627],
+            [-0.01436, -0.02519, -0.01887],
+            [0.13145, -0.0941, 0.02256],
+            [-0.0128, -0.00109, 0.01468],
+            [0.03173, -0.01365, 0.01746],
+            [-0.05829, 0.01014, -0.02845],
+            [-0.02415, 0.01729, 0.03453],
+            [0.0132, 0.03026, 0.00655],
+            [0.004, 0.00867, -0.01574],
+            [0.07235, -0.00765, 0.03589],
+            [0.05798, 0.01141, 0.00962],
+            [-0.01311, -0.02256, -0.04102],
+            [-0.03182, -0.02579, 0.00077],
+            [0.02911, 0.01766, -0.01212],
+            [0.01937, -0.00168, -0.02601],
+        ],
+    },
 }

extensions_built_in/inference_engine/latent_preview.py

@@ -50,6 +50,7 @@ ARCH_LATENT_FORMAT = {
     "anima": "Wan21",
     "cogview4": "Flux",
     "prx_pixel": None,  # pixel space
+    "flux3": "Flux3Image",
     "iris3b": None,  # pixel space
 }
 

toolkit/models/registry.py

@@ -55,6 +55,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "wan22_14b_i2v": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}, "needs_control_image": True},
     "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
+    "flux3": {"modality": "image", "model": {"name_or_path": "black-forest-labs/flux-3-image-klein-base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}},
     "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}},
     "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}},
     "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},