Add support for Bria Fibo 2.0

README.md

@@ -42,6 +42,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow)
 - [OpenMOSS-Team/MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) (MOSS-Image SFT)
 - [OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) (MOSS-Image Base)
+- [briaai/Fibo-2.0-dev](https://huggingface.co/briaai/Fibo-2.0-dev) (Fibo 2.0 Dev)
 
 ### Instruction / Edit
 - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev)

extensions_built_in/diffusion_models/__init__.py

@@ -25,6 +25,7 @@ from .mageflow import MageFlowModel, MageFlowEditModel
 from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel, MinimaxH3OmniModel
 from .flux3 import Flux3Model
 from .moss_image import MossImageModel
+from .fibo2 import Fibo2Model
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
@@ -70,4 +71,5 @@ AI_TOOLKIT_MODELS = [
     MinimaxH3OmniModel,
     Flux3Model,
     MossImageModel,
+    Fibo2Model,
 ]

extensions_built_in/diffusion_models/fibo2/__init__.py

@@ -0,0 +1,3 @@
+from .fibo2 import Fibo2Model
+
+__all__ = ["Fibo2Model"]

extensions_built_in/diffusion_models/fibo2/fibo2.py

@@ -0,0 +1,467 @@
+"""Fibo 2.0 dev (briaai/Fibo-2.0-dev) for ai-toolkit, text-to-image and edit.
+
+  - DiT: 44 Z-Image blocks, 3584 wide, a Perceiver turning the text into
+    256-512 gist tokens and five cross-attention injection blocks
+    (``src/transformer.py``, Comfy-Org/Fibo-2.0 key layout)
+  - VAE: FLUX.2, 2x2 packed and BatchNorm normalized (the toolkit's BFL
+    ``AutoEncoder``, the shared comfy flux2-vae file)
+  - text: Qwen3-VL-4B, six RMS-normalized bundles of three hidden layers;
+    stored per item as (L, 46080) in an ``AdvancedPromptEmbeds``
+
+Editing is not a separate model: control images in the dataset (or sample
+ctrl_img) make it an edit. References ride into the prompt through Qwen3-VL's
+vision tower and into the DiT as clean t = 0 latents; without them it is
+text-to-image. model_kwargs.match_target_res (default true) sizes references
+against the target, see ``src/references.py``.
+
+``model.name_or_path`` defaults to the Comfy-Org/Fibo-2.0 repack (int8 convrot
+for the default convrot8 qtype, bf16 otherwise); a comfy-layout
+``.safetensors`` (local or ``org/repo/file``) loads as named. Prompts are
+structured JSON captions (cleaned and compacted as in training) or plain text.
+
+Flow convention matches ai-toolkit (sigma 1 = noise, target = noise - clean).
+"""
+
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+import torchvision.transforms.functional as TF
+from PIL import Image
+from safetensors.torch import load_file
+from transformers import AutoImageProcessor
+
+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.models.v2.vae.flux2_kl import AutoEncoder, convert_diffusers_state_dict
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.checkpoints import (
+    QWEN3_VL_4B_REPO,
+    comfy_text_encoder_file,
+    comfy_transformer_file,
+    comfy_vae_file,
+    resolve_hub_file,
+)
+from .src.pipeline import Fibo2Pipeline, scheduler_config, stack_text_embeds
+from .src.references import drop_repeats, prepare_references, vlm_image
+from .src.text_encoder import MAX_TEXT_TOKENS, Fibo2TextEncoder, encode_prompts
+from .src.transformer import Fibo2Transformer
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+
+class Fibo2Model(BaseModel):
+    arch = "fibo2"
+
+    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 = ["Fibo2Transformer"]
+        self.supports_model_paths = True
+        self.use_old_lokr_format = False
+        self.lora_keys_use_comfy_prefix = True
+        self.max_text_length = int(
+            self.model_config.model_kwargs.get("max_text_length", MAX_TEXT_TOKENS)
+        )
+        # edit when the dataset has control images, text-to-image otherwise
+        self.encode_control_in_text_embeddings = True
+        self.has_multiple_control_images = True
+        # references are sized in prepare_references, identically for the
+        # cached prompt and the training latents
+        self.use_raw_control_images = True
+        self.image_processor = None
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    def get_bucket_divisibility(self):
+        return 16  # f8 conv + 2x2 packing
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def _load_transformer(self) -> Fibo2Transformer:
+        self.print_and_status_update("Loading transformer (Fibo2Transformer)")
+        qtype = self.model_config.qtype if self.model_config.quantize else None
+        path = comfy_transformer_file(
+            self.model_config.name_or_path,
+            qtype=qtype,
+            hf_token=HF_TOKEN,
+            status_fn=self.print_and_status_update,
+        )
+        self.print_and_status_update(f" - transformer weights: {path}")
+        return Fibo2Transformer.load_model(
+            path, dtype=self.torch_dtype, qtype=qtype, quantize_on_load=False
+        )
+
+    def _load_text_encoder(self) -> Fibo2TextEncoder:
+        te_kwargs = self.component_load_kwargs("te")
+        te_path = self.model_config.te_name_or_path
+        if te_path is None:
+            te_path = comfy_text_encoder_file(
+                qtype=te_kwargs["qtype"],
+                hf_token=HF_TOKEN,
+                status_fn=self.print_and_status_update,
+            )
+        elif te_path.endswith(".safetensors"):
+            te_path = resolve_hub_file(
+                te_path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}")
+        text_encoder = Fibo2TextEncoder.load_model(
+            te_path,
+            dtype=te_kwargs["dtype"],
+            config_path=QWEN3_VL_4B_REPO if te_path.endswith(".safetensors") else None,
+            subfolder="",
+            qtype=te_kwargs["qtype"],
+            quantize_on_load=False,
+            use_comfy_weights=False,
+        )
+        text_encoder.patch_vision_patch_embed()
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        return text_encoder
+
+    def _load_vae(self) -> AutoEncoder:
+        self.print_and_status_update("Loading VAE")
+        path = self.model_config.vae_path
+        if path is None:
+            path = comfy_vae_file(
+                hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        elif not os.path.exists(path):
+            path = resolve_hub_file(
+                path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        state_dict = load_file(path)
+        if any(k.startswith("encoder.down_blocks.") for k in state_dict):
+            state_dict = convert_diffusers_state_dict(state_dict)
+        vae = AutoEncoder.load_from_state_dict(state_dict, self.vae_torch_dtype)
+        return vae.eval().requires_grad_(False)
+
+    def load_model(self):
+        self.print_and_status_update("Loading Fibo 2.0 model")
+
+        transformer = self._load_transformer()
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
+        flush()
+
+        text_encoder = self._load_text_encoder()
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        flush()
+        tokenizer = Fibo2TextEncoder.load_tokenizer()
+        self.image_processor = Fibo2TextEncoder._local_first(
+            AutoImageProcessor.from_pretrained, QWEN3_VL_4B_REPO
+        )
+
+        vae = self._load_vae()
+        vae.to(self.vae_device_torch)
+
+        self.noise_scheduler = Fibo2Model.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.model = transformer
+        self.pipeline = Fibo2Pipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Generation (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return Fibo2Pipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: Fibo2Pipeline,
+        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)
+
+        # the same references the sampler encoded into the prompt
+        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 paths:
+                paths.append(path)
+        ref_latents = None
+        if paths:
+            crops = prepare_references(
+                [TF.to_tensor(Image.open(p).convert("RGB")) for p in paths],
+                (gen_config.width, gen_config.height),
+                self.match_target_res,
+            )
+            ref_latents = self._encode_crops(crops)
+            # the vendor's negative is text only (the references still
+            # condition both passes); the sampler encoded it with them
+            if gen_config.negative_prompt.strip() and self.sample_prompts_cache is None:
+                unconditional_embeds = self.get_prompt_embeds(
+                    gen_config.negative_prompt
+                )
+
+        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,
+            ref_latents=ref_latents,
+        )
+        return result[0]
+
+    # ------------------------------------------------------------------
+    # References (edit)
+    # ------------------------------------------------------------------
+    @property
+    def match_target_res(self) -> bool:
+        return bool(self.model_config.model_kwargs.get("match_target_res", True))
+
+    @property
+    def text_embedding_uses_target_size(self) -> bool:
+        # the dataloader adds the item's bucket size to the text-embedding cache key
+        return self.match_target_res
+
+    def get_text_embedding_space_version(self) -> str:
+        # reference sizing changes the vision tokens; keep caches from different rules apart
+        rule = "match" if self.match_target_res else "trained"
+        return f"{self.text_embedding_space_version}_ref{rule}"
+
+    def _target_size(self, target_size):
+        if target_size is None:
+            return None
+        sc = self.get_bucket_divisibility()
+        return (int(target_size[0] // sc * sc), int(target_size[1] // sc * sc))
+
+    @staticmethod
+    def _normalize_control_images(
+        control_images, batch_size: int
+    ) -> List[List[torch.Tensor]]:
+        """A (B, C, H, W) / (B, N, C, H, W) tensor, a per-sample list of lists
+        or a flat list for one prompt -> one list of [0, 1] tensors per item."""
+        if control_images is None:
+            return [[] for _ in range(batch_size)]
+        if isinstance(control_images, torch.Tensor):
+            if control_images.dim() == 5:
+                control_images = [list(sample) for sample in control_images]
+            else:
+                control_images = [[sample] for sample in control_images]
+        elif len(control_images) > 0 and not isinstance(control_images[0], list):
+            control_images = [list(control_images)]
+        control_images = [drop_repeats(sample) for sample in control_images]
+        if len(control_images) == 1 and batch_size > 1:
+            control_images = control_images * batch_size
+        if len(control_images) != batch_size:
+            raise ValueError(
+                f"got {len(control_images)} control image sets for {batch_size} prompts"
+            )
+        return control_images
+
+    def _encode_crops(self, crops) -> List[torch.Tensor]:
+        """Reference crops -> one (1, 128, h, w) latent each."""
+        return [
+            self.encode_images(
+                [TF.to_tensor(crop) * 2 - 1],
+                device=self.device_torch,
+                dtype=self.torch_dtype,
+            )
+            for crop in crops
+        ]
+
+    def _batch_ref_latents(
+        self, batch, batch_size: int, target_size
+    ) -> Optional[List[torch.Tensor]]:
+        control = batch.control_tensor_list
+        if control is None:
+            control = batch.control_tensor
+        if control is None:
+            return None
+        samples = [
+            self._encode_crops(
+                prepare_references(images, target_size, self.match_target_res)
+            )
+            for images in self._normalize_control_images(control, batch_size)
+        ]
+        if not any(samples):
+            return None
+        shapes = {tuple(tuple(lat.shape) for lat in sample) for sample in samples}
+        if len(shapes) > 1:
+            raise ValueError(
+                "every sample in a batch needs the same number and size of references; "
+                f"got {sorted(shapes)}. Use match_target_res or batch size 1."
+            )
+        return [torch.cat(refs, 0) for refs in zip(*samples)]
+
+    # ------------------------------------------------------------------
+    # 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]
+        embeds_list = text_embeddings.text_embeds
+        if len(embeds_list) == 1 and batch_size > 1:
+            embeds_list = embeds_list * batch_size
+        with torch.no_grad():
+            embeds, mask = stack_text_embeds(embeds_list, device, dtype)
+            t = timestep.to(device, dtype=torch.float32).reshape(-1)
+            if t.shape[0] != batch_size:
+                t = t.expand(batch_size)
+            ref_latents = None
+            if batch is not None:
+                # the bucket the prompt was cached against
+                _, _, h, w = latent_model_input.shape
+                ref_latents = self._batch_ref_latents(
+                    batch, batch_size, (w * 16, h * 16)
+                )
+
+        return self.transformer(
+            latent_model_input.to(device, dtype), t, embeds, mask, ref_latents
+        )
+
+    @torch.no_grad()
+    def get_prompt_embeds(
+        self, prompt, control_images=None, target_size=None
+    ) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+        target_size = self._target_size(target_size)
+        images = [
+            [
+                vlm_image(crop)
+                for crop in prepare_references(
+                    sample, target_size, self.match_target_res
+                )
+            ]
+            for sample in self._normalize_control_images(control_images, len(prompt))
+        ]
+        embeds = encode_prompts(
+            self.text_encoder,
+            self.tokenizer,
+            self.image_processor,
+            prompt,
+            images,
+            self.device_torch,
+            self.torch_dtype,
+            max_length=self.max_text_length,
+        )
+        return AdvancedPromptEmbeds(text_embeds=embeds)
+
+    def condition_noisy_latents(
+        self, latents: torch.Tensor, batch: "DataLoaderBatchDTO"
+    ):
+        # references join the sequence in get_noise_prediction, clean
+        return latents.detach()
+
+    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] -> normalized latents (B, 128, H/16, W/16)."""
+        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)
+        batch = torch.stack(
+            [img.to(device, dtype=self.vae.dtype) for img in image_list]
+        )
+        return self.vae.encode(batch).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):
+        if not output_path.endswith(".safetensors"):
+            output_path = output_path + ".safetensors"
+        transformer: Fibo2Transformer = unwrap_model(self.model)
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
+            quantized=False,
+        )
+
+    def get_base_model_version(self):
+        return "fibo2"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return Fibo2Transformer.get_transformer_block_names()
+
+    def get_quantization_exclude_modules(self):
+        return Fibo2Transformer.get_quantization_exclude_modules()

extensions_built_in/diffusion_models/fibo2/src/checkpoints.py

@@ -0,0 +1,100 @@
+"""Where Fibo 2.0 weights come from.
+
+Weights live in the ComfyUI layout under MODELS_PATH, named like the
+Comfy-Org/Fibo-2.0 repack (the vendor repo briaai/Fibo-2.0-dev is gated and
+diffusers-keyed; not a source here).
+
+- ``Comfy-Org/Fibo-2.0`` (the UI default): the dev DiT, int8 convrot or bf16,
+  whichever is local (ranked for the requested qtype), else the best-ranked
+  one is downloaded.
+- ``org/repo/path/file.safetensors`` or a local file: that file (fine-tunes
+  saved by the toolkit are in the same layout).
+
+The text encoder is stock Qwen3-VL-4B-Instruct (the repack's bf16 is byte
+identical to Comfy-Org/Qwen3-VL's), so config + tokenizer come from the Qwen
+repo. The VAE is the shared FLUX.2 file.
+"""
+
+import os
+from typing import Callable, Optional
+
+from toolkit.models.v2.resolver import comfy_precision_rank, resolve_comfy_candidates
+
+COMFY_REPO = "Comfy-Org/Fibo-2.0"
+QWEN3_VL_4B_REPO = "Qwen/Qwen3-VL-4B-Instruct"
+
+COMFY_TRANSFORMER_FILES = [
+    "diffusion_models/fibo2_dev_int8_convrot.safetensors",
+    "diffusion_models/fibo2_dev_bf16.safetensors",
+]
+COMFY_TEXT_ENCODER_FILES = [
+    "text_encoders/qwen3vl_4b_int8_convrot.safetensors",
+    "text_encoders/qwen3vl_4b_bf16.safetensors",
+]
+COMFY_VAE_FILE = "vae/flux2-vae.safetensors"
+
+
+def resolve_hub_file(
+    ref: str,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    """A local path, or 'org/repo/rel/file.safetensors' found at / downloaded
+    to its comfy spot under MODELS_PATH."""
+    if os.path.isfile(ref):
+        return ref
+    parts = ref.split("/")
+    if len(parts) < 3:
+        raise FileNotFoundError(
+            f"{ref} is not a local file or an 'org/repo/path/file.safetensors' reference"
+        )
+    return resolve_comfy_candidates(
+        ["/".join(parts[2:])],
+        repo_id="/".join(parts[:2]),
+        hf_token=hf_token,
+        status_fn=status_fn,
+    )
+
+
+def _best_comfy_file(candidates, qtype, hf_token, status_fn) -> str:
+    found = resolve_comfy_candidates(
+        candidates, repo_id=COMFY_REPO, local_only=True, qtype=qtype
+    )
+    if found is not None:
+        return found
+    best = min(candidates, key=lambda c: comfy_precision_rank(c, qtype=qtype))
+    return resolve_comfy_candidates(
+        [best], repo_id=COMFY_REPO, hf_token=hf_token, status_fn=status_fn
+    )
+
+
+def comfy_transformer_file(
+    name_or_path: str,
+    qtype: Optional[str] = None,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    if name_or_path.endswith(".safetensors"):
+        return resolve_hub_file(name_or_path, hf_token=hf_token, status_fn=status_fn)
+    if name_or_path != COMFY_REPO:
+        raise ValueError(
+            f"Fibo 2.0 loads {COMFY_REPO} or a comfy-layout .safetensors file, not {name_or_path}"
+        )
+    return _best_comfy_file(COMFY_TRANSFORMER_FILES, qtype, hf_token, status_fn)
+
+
+def comfy_text_encoder_file(
+    qtype: Optional[str] = None,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    return _best_comfy_file(COMFY_TEXT_ENCODER_FILES, qtype, hf_token, status_fn)
+
+
+def comfy_vae_file(
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    return resolve_comfy_candidates(
+        [COMFY_VAE_FILE], repo_id=COMFY_REPO, hf_token=hf_token, status_fn=status_fn
+    )

extensions_built_in/diffusion_models/fibo2/src/pipeline.py

@@ -0,0 +1,138 @@
+"""Fibo 2.0 sampling (training previews / inference engine).
+
+Matches diffusers' BriaFibo2Pipeline: sigmas linspace(1, 1 / steps) under
+FLUX.2's resolution-dependent exponential shift (mu 0.5 at 256 tokens to 1.15
+at 4096), fp32 Euler on fp32 noise, CFG 5 against the null caption as a
+separate pass. Edit references condition both passes as clean latents. Flow
+convention is the toolkit's: sigma 1 = noise, the DiT predicts noise - clean.
+"""
+
+import math
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+DEFAULT_GUIDANCE = 5.0
+DEFAULT_STEPS = 30
+NUM_TRAIN_TIMESTEPS = 1000
+
+# FLUX.2's scheduler config, which Fibo 2.0 shares
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": 0.5,
+    "max_image_seq_len": 4096,
+    "max_shift": 1.15,
+    "num_train_timesteps": NUM_TRAIN_TIMESTEPS,
+    "shift": 3.0,
+    "use_dynamic_shifting": True,
+}
+
+
+def fibo2_sigmas(num_steps: int, image_seq_len: int) -> torch.Tensor:
+    """Shifted sigmas with the terminal 0 appended."""
+    c = scheduler_config
+    m = (c["max_shift"] - c["base_shift"]) / (
+        c["max_image_seq_len"] - c["base_image_seq_len"]
+    )
+    mu = image_seq_len * m + c["base_shift"] - m * c["base_image_seq_len"]
+    sigmas = np.linspace(1.0, 1 / num_steps, num_steps)
+    sigmas = math.exp(mu) / (math.exp(mu) + (1 / sigmas - 1))
+    return torch.cat([torch.from_numpy(sigmas).float(), torch.zeros(1)])
+
+
+def stack_text_embeds(
+    embeds: List[torch.Tensor], device: torch.device, dtype: torch.dtype
+):
+    """Per-prompt (L_i, D) bundles -> right-padded (B, S, D) + (B, S) mask,
+    the mask None when nothing is padded."""
+    length = max(e.shape[0] for e in embeds)
+    stacked = torch.stack(
+        [F.pad(e.to(device, dtype), (0, 0, 0, length - e.shape[0])) for e in embeds]
+    )
+    if all(e.shape[0] == length for e in embeds):
+        return stacked, None
+    mask = torch.stack(
+        [torch.arange(length, device=device) < e.shape[0] for e in embeds]
+    )
+    return stacked, mask
+
+
+class Fibo2Pipeline:
+    """Lightweight sampler used by ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # the Fibo2Model, for its transformer + decode
+        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 = DEFAULT_STEPS,
+        guidance_scale: float = DEFAULT_GUIDANCE,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        ref_latents: Optional[List[torch.Tensor]] = None,  # each (1, 128, h, w)
+        **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
+        cond, cond_mask = stack_text_embeds(
+            conditional_embeds.text_embeds, device, dtype
+        )
+        if do_cfg:
+            uncond, uncond_mask = stack_text_embeds(
+                unconditional_embeds.text_embeds, device, dtype
+            )
+
+        shape = (1, transformer.params.in_channels, height // 16, width // 16)
+        if latents is None:
+            latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=torch.float32
+            )
+        latents = latents.to(device, torch.float32)
+
+        sigmas = fibo2_sigmas(num_inference_steps, shape[2] * shape[3]).to(device)
+        for step_index in range(num_inference_steps):
+            sigma, sigma_next = sigmas[step_index], sigmas[step_index + 1]
+            timestep = (sigma * NUM_TRAIN_TIMESTEPS).reshape(1)
+            x = latents.to(dtype)
+            pred = transformer(x, timestep, cond, cond_mask, ref_latents).float()
+            if do_cfg:
+                pred_uncond = transformer(
+                    x, timestep, uncond, uncond_mask, ref_latents
+                ).float()
+                pred = pred_uncond + guidance_scale * (pred - pred_uncond)
+            if getattr(model, "sample_step_hook", None) is not None:
+                # inference engine preview: x0 estimate
+                model._emit_sample_step(
+                    latents - sigma * pred, step_index, num_inference_steps
+                )
+            latents = latents + (sigma_next - sigma) * pred
+
+        images = model.decode_latents(latents, device=device, dtype=torch.float32)
+        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/fibo2/src/references.py

@@ -0,0 +1,129 @@
+"""Fibo 2.0 edit references: sizing, crops, and Qwen3-VL's copies.
+
+Bria's preprocessing (diffusers BriaFibo2ImageProcessor): a reference is
+scaled to cover its size (Lanczos, ceil) and center cropped; the VAE encodes
+the crop, Qwen3-VL sees it thumbnailed to at most 768 px a side on the
+processor's 32 px grid. Vendor sizes: a single reference crops to the output
+size (the edit lines up with it pixel for pixel), several references each to
+the ~1 MP size Fibo 2.0 was trained at for their own aspect ratio.
+
+match_target_res (default on): a single reference crops to the target size,
+several each to the target's area at their own aspect. Off: every reference
+at its own trained ~1 MP size, whatever the target.
+"""
+
+import math
+from functools import lru_cache
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+from PIL import Image
+
+TRAINED_AREA = 1024 * 1024
+VLM_MAX_SIDE = 768
+# Qwen3-VL processor grid for Fibo 2.0 (ComfyUI's port of Bria's processor config)
+VLM_GRID = 32
+VLM_MIN_PIXELS = 3136
+VLM_MAX_PIXELS = 1605632
+
+
+@lru_cache(maxsize=256)
+def trained_size(aspect_ratio: float, area: int = TRAINED_AREA) -> Tuple[int, int]:
+    """(width, height): of the multiples of 16 (>= 64 a side) within 10% of
+    area, the least cropping for this aspect, then the closest to area."""
+    sizes = [
+        (w, h)
+        for w in range(64, 2 * area // 64, 16)
+        for h in range(
+            max(64, 16 * math.ceil(0.9 * area / w / 16)),
+            16 * math.floor(1.1 * area / w / 16) + 1,
+            16,
+        )
+    ]
+    return min(
+        sizes,
+        key=lambda s: (
+            round(1 - min(s[0] / s[1] / aspect_ratio, aspect_ratio * s[1] / s[0]), 9),
+            abs(s[0] * s[1] - area),
+            -s[0],
+        ),
+    )
+
+
+def reference_sizes(
+    images: Sequence[Image.Image],
+    target_size: Optional[Tuple[int, int]],
+    match_target_res: bool,
+) -> List[Tuple[int, int]]:
+    """(width, height) each reference is cropped to."""
+    if match_target_res and target_size is not None:
+        width, height = target_size
+        if len(images) == 1:
+            return [(width, height)]
+        return [trained_size(img.width / img.height, width * height) for img in images]
+    return [trained_size(img.width / img.height) for img in images]
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """(C, H, W) / (1, C, H, W) in [0, 1] -> RGB PIL."""
+    if image.dim() == 4:
+        image = image[0]
+    array = (image[:3].float().clamp(0, 1) * 255).round().to(torch.uint8)
+    return Image.fromarray(array.permute(1, 2, 0).cpu().numpy())
+
+
+def resize_and_crop(image: Image.Image, width: int, height: int) -> Image.Image:
+    scale = max(height / image.height, width / image.width)
+    resized_width, resized_height = (
+        math.ceil(image.width * scale),
+        math.ceil(image.height * scale),
+    )
+    image = image.resize((resized_width, resized_height), resample=Image.LANCZOS)
+    left, top = (resized_width - width) // 2, (resized_height - height) // 2
+    return image.crop((left, top, left + width, top + height))
+
+
+def vlm_image(crop: Image.Image) -> Image.Image:
+    """Qwen3-VL's copy of a reference crop: thumbnailed, then on the 32 px
+    grid (pixel budget clamped like the processor's smart resize)."""
+    image = crop.copy()
+    image.thumbnail((VLM_MAX_SIDE, VLM_MAX_SIDE))
+    h, w = image.height, image.width
+    height, width = round(h / VLM_GRID) * VLM_GRID, round(w / VLM_GRID) * VLM_GRID
+    if height * width > VLM_MAX_PIXELS:
+        beta = math.sqrt(h * w / VLM_MAX_PIXELS)
+        height = max(VLM_GRID, math.floor(h / beta / VLM_GRID) * VLM_GRID)
+        width = max(VLM_GRID, math.floor(w / beta / VLM_GRID) * VLM_GRID)
+    elif height * width < VLM_MIN_PIXELS:
+        beta = math.sqrt(VLM_MIN_PIXELS / (h * w))
+        height = math.ceil(h * beta / VLM_GRID) * VLM_GRID
+        width = math.ceil(w * beta / VLM_GRID) * VLM_GRID
+    if (width, height) != (w, h):
+        image = image.resize((width, height), resample=Image.BICUBIC)
+    return image
+
+
+def prepare_references(
+    images: Sequence[torch.Tensor],
+    target_size: Optional[Tuple[int, int]],
+    match_target_res: bool,
+) -> List[Image.Image]:
+    """[0, 1] reference tensors -> the crops the VAE and Qwen3-VL see."""
+    pils = [tensor_to_pil(image) for image in images]
+    sizes = reference_sizes(pils, target_size, match_target_res)
+    return [resize_and_crop(img, w, h) for img, (w, h) in zip(pils, sizes)]
+
+
+def drop_repeats(images):
+    """ctrl_img_1 defaults to ctrl_img, so a single reference can arrive twice."""
+    kept = []
+    for image in images:
+        previous = kept[-1] if kept else None
+        if (
+            previous is not None
+            and previous.shape == image.shape
+            and torch.equal(previous, image)
+        ):
+            continue
+        kept.append(image)
+    return kept

extensions_built_in/diffusion_models/fibo2/src/text_encoder.py

@@ -0,0 +1,227 @@
+"""Fibo 2.0 text conditioning: Qwen3-VL-4B hidden states in six bundles.
+
+The prompt goes in Qwen3-VL's user turn; 15 hidden layers are each scaled to
+unit RMS in fp32, and every bundle concatenates three of them: the
+Perceiver's, then one per injection block. Stored per prompt as
+(L, 6 * 7680), the ComfyUI conditioning layout.
+
+- Generation: a structured JSON caption, cleaned and compacted the way Fibo
+  2.0 was trained, or plain text as it is; an empty prompt is the null
+  caption (CFG negative, caption dropout).
+- Edit: one JSON object, a vision marker per reference (``image`` or
+  ``image_1``..), the result's caption when the prompt has one, then the edit
+  instruction; Qwen3-VL sees each reference in place of its marker.
+
+Not done: the images of a caption's text_render entries (diffusers' LoRA
+trainers encode without them).
+"""
+
+import json
+import math
+from typing import List, Optional, Sequence
+
+import torch
+from PIL import Image
+
+from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLModelEncoder
+
+from .checkpoints import QWEN3_VL_4B_REPO
+
+PROMPT_TEMPLATE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+IMAGE_PAD = "<|image_pad|>"
+VISION_MARKER = f"<|vision_start|>{IMAGE_PAD}<|vision_end|>"
+
+# the null caption Fibo 2.0 was trained with for classifier-free guidance
+NULL_CAPTION = (
+    '{"short_description":"","objects":[],"background_setting":"","lighting":{},'
+    '"aesthetics":{},"photographic_characteristics":{},"style_medium":"",'
+    '"context":"","artistic_style":""}'
+)
+
+# Qwen3-VL hidden_states indices (0 = embeddings) behind the six bundles
+BUNDLE_LAYERS = (
+    (9, 20, 31),
+    (5, 16, 27),
+    (7, 18, 29),
+    (9, 20, 31),
+    (11, 22, 33),
+    (13, 24, 35),
+)
+MAX_TEXT_TOKENS = 4096
+SCORES = ("aesthetic_score", "preference_score")
+
+
+def _dumps(value) -> str:
+    return json.dumps(value, separators=(",", ":"), ensure_ascii=False)
+
+
+def _empty(value) -> bool:
+    if isinstance(value, float):
+        return math.isnan(value)
+    return value is None or (isinstance(value, (str, dict, list)) and not value)
+
+
+def _clean(value):
+    if isinstance(value, dict):
+        value = {k: _clean(v) for k, v in value.items()}
+        return {k: v for k, v in value.items() if not _empty(v)}
+    if isinstance(value, list):
+        return [v for v in map(_clean, value) if not _empty(v)]
+    return value
+
+
+def parse_prompt(prompt: str):
+    """(caption, edit instruction) as diffusers' BriaFibo2Pipeline parses
+    them; a prompt that isn't a JSON object comes back as it is."""
+    try:
+        caption = json.loads(prompt)
+    except json.JSONDecodeError:
+        return prompt, None
+    if not isinstance(caption, dict):
+        return prompt, None
+    instruction = caption.pop("edit_instruction", None)
+    if "structured_caption" in caption:
+        caption = caption.pop("structured_caption")
+        if not isinstance(caption, dict):
+            return caption or None, instruction
+    keys = set(caption)
+    caption = {
+        key: {k: v for k, v in value.items() if k not in SCORES}
+        if isinstance(value, dict)
+        else value
+        for key, value in _clean(caption).items()
+        if key not in SCORES
+    }
+    caption = {key: value for key, value in caption.items() if value != {}}
+    # exactly a short_description and a text_render list keeps both, the list
+    # even when empty: the image has no text
+    if keys == {"short_description", "text_render"}:
+        caption = {
+            "short_description": caption.get("short_description") or "",
+            "text_render": caption.get("text_render") or [],
+        }
+    return caption or None, instruction
+
+
+def generation_text(prompt: str) -> str:
+    if prompt.strip() == "":
+        return NULL_CAPTION
+    caption, _ = parse_prompt(prompt)
+    return prompt if isinstance(caption, str) else _dumps(caption or {})
+
+
+def edit_text(prompt: str, num_images: int) -> str:
+    caption, instruction = parse_prompt(prompt)
+    if instruction is None:
+        caption, instruction = None, prompt
+    if num_images == 1:
+        edit = {"image": VISION_MARKER}
+    else:
+        edit = {f"image_{k}": VISION_MARKER for k in range(1, num_images + 1)}
+    if caption:
+        edit["structured_caption"] = caption
+    edit["edit_instruction"] = instruction
+    return _dumps(edit)
+
+
+class Fibo2TextEncoder(Qwen3VLModelEncoder):
+    """Qwen3-VL-4B (language + vision tower) from the Comfy single file, int8
+    convrot or bf16; the tied lm_head is dropped on load."""
+
+    aitk_config_repo = QWEN3_VL_4B_REPO
+    aitk_tokenizer_repo = QWEN3_VL_4B_REPO
+    aitk_tokenizer_subfolder = ""
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        return {
+            key[len("model.") :] if key.startswith("model.") else key: value
+            for key, value in state_dict.items()
+            if not key.startswith("lm_head.")
+        }
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        return {f"model.{k}": v for k, v in state_dict.items()}
+
+
+def _bundles(hidden_states, dtype) -> torch.Tensor:
+    normalized = {}
+    for layer in sorted({layer for bundle in BUNDLE_LAYERS for layer in bundle}):
+        states = hidden_states[layer].float()
+        normalized[layer] = (
+            states * torch.rsqrt(states.pow(2).mean(-1, keepdim=True) + 1e-6)
+        ).to(dtype)
+    return torch.cat(
+        [normalized[layer] for bundle in BUNDLE_LAYERS for layer in bundle], -1
+    )
+
+
+@torch.no_grad()
+def encode_prompts(
+    text_encoder,
+    tokenizer,
+    image_processor,
+    prompts: List[str],
+    images: Sequence[Optional[List[Image.Image]]],
+    device: torch.device,
+    dtype: torch.dtype,
+    max_length: int = MAX_TEXT_TOKENS,
+) -> List[torch.Tensor]:
+    """Prompts (+ per-prompt Qwen3-VL reference images) -> per-prompt
+    (L_i, 46080) bundles. Text-only prompts share one right-padded forward
+    without a mask: causal attention keeps the pads from reaching any real
+    token. Prompts with references run one forward each."""
+    out: List[Optional[torch.Tensor]] = [None] * len(prompts)
+
+    text_only = [
+        i for i, imgs in enumerate(images) if not imgs or prompts[i].strip() == ""
+    ]
+    if text_only:
+        enc = tokenizer(
+            [PROMPT_TEMPLATE.format(generation_text(prompts[i])) for i in text_only],
+            return_tensors="pt",
+            padding=True,
+            padding_side="right",
+            truncation=True,
+            max_length=max_length,
+        )
+        lengths = enc["attention_mask"].sum(1).tolist()
+        hidden_states = text_encoder(
+            input_ids=enc["input_ids"].to(device),
+            output_hidden_states=True,
+            use_cache=False,
+        ).hidden_states
+        context = _bundles(hidden_states, dtype)
+        for row, (i, n) in enumerate(zip(text_only, lengths)):
+            out[i] = context[row, :n]
+
+    image_token_id = text_encoder.config.image_token_id
+    merge = image_processor.merge_size**2
+    for i, imgs in enumerate(images):
+        if out[i] is not None:
+            continue
+        pixels = image_processor(images=imgs, do_resize=False, return_tensors="pt")
+        grid = pixels["image_grid_thw"]
+        text = PROMPT_TEMPLATE.format(edit_text(prompts[i], len(imgs)))
+        # each marker's single pad expands to its image's token count; no
+        # truncation, the pad runs must stay whole
+        parts = text.split(IMAGE_PAD)
+        counts = (grid.prod(-1) // merge).tolist()
+        text = parts[0] + "".join(
+            IMAGE_PAD * n + part for n, part in zip(counts, parts[1:])
+        )
+        input_ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
+        mm_token_type_ids = (input_ids == image_token_id).int()
+        if int(mm_token_type_ids.sum()) != sum(counts):
+            raise ValueError("image placeholder count does not match the image grids")
+        hidden_states = text_encoder(
+            input_ids=input_ids,
+            pixel_values=pixels["pixel_values"].to(device, dtype),
+            image_grid_thw=grid.to(device),
+            mm_token_type_ids=mm_token_type_ids,
+            output_hidden_states=True,
+            use_cache=False,
+        ).hidden_states
+        out[i] = _bundles(hidden_states, dtype)[0]
+    return out

extensions_built_in/diffusion_models/fibo2/src/transformer.py

@@ -0,0 +1,479 @@
+"""Fibo 2.0 DiT (briaai/Fibo-2.0-dev).
+
+Port of diffusers' ``BriaFibo2Transformer2DModel`` in the Comfy-Org/Fibo-2.0
+key layout: self-attention q/k/v fused into ``attention.qkv``, cross-attention
+and Perceiver k/v into ``to_kv``, SwiGLU w1/w3 into ``feed_forward.w13`` with
+the hidden width zero-padded to a multiple of 256 (9536 -> 9728), and the
+input taking FLUX.2's 128-channel packed latents (the diffusers model's 2x2
+patchify of the 32-channel latents, same (c, dy, dx) feature order). Comfy
+files load as-is and LoRA keys match ComfyUI.
+
+One stream of [gist | image | references] through Z-Image blocks; the gist
+tokens come from a Perceiver over the first Qwen3-VL bundle and are clean
+(t = 0 modulation) like the references. Five injection blocks also
+cross-attend to their own bundle. Batches whose gist counts differ put the
+gist last and mask the padding, exactly as diffusers does.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+
+ADALN_DIM = 256
+
+
+@dataclass
+class Fibo2Params:
+    dim: int = 3584
+    n_heads: int = 28
+    n_layers: int = 44
+    n_refiner_layers: int = 2
+    injection_layers: Tuple[int, ...] = (6, 14, 22, 30, 38)
+    cap_feat_dim: int = 7680
+    perceiver_layers: int = 2
+    max_gist_tokens: int = 512
+    ffn_dim: int = 9728
+    in_channels: int = 128
+    norm_eps: float = 1e-5
+    # fixed by the release (diffusers config defaults)
+    min_gist_tokens: int = 256
+    gist_step: int = 64
+    gist_min_text_len: int = 256
+    gist_max_text_len: int = 2048
+    rope_theta: float = 256.0
+    axes_dims: Tuple[int, ...] = (32, 48, 48)
+
+
+def gist_counts(lengths: Sequence[int], p: Fibo2Params) -> List[int]:
+    """Gist tokens per prompt: min_gist up to min_text_len text tokens, rising
+    linearly to max_gist at max_text_len, rounded up to gist_step."""
+    span = p.gist_max_text_len - p.gist_min_text_len
+    counts = []
+    for n in lengths:
+        ratio = min(1.0, max(0.0, (float(n) - p.gist_min_text_len) / span))
+        raw = p.min_gist_tokens + ratio * (p.max_gist_tokens - p.min_gist_tokens)
+        quantized = int(math.ceil(raw / p.gist_step)) * p.gist_step
+        counts.append(max(p.min_gist_tokens, min(p.max_gist_tokens, quantized)))
+    return counts
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = F.rms_norm(x.float(), (x.shape[-1],), eps=self.eps)
+        return y.to(x.dtype) * self.weight
+
+
+def _apply_rope(x, rope):
+    # interleaved (real, imag) pairs rotated in fp32, as diffusers' complex form
+    cos, sin = rope
+    re, im = x.float().unflatten(-1, (-1, 2)).unbind(-1)
+    out = torch.stack((re * cos - im * sin, re * sin + im * cos), -1).flatten(-2)
+    return out.type_as(x)
+
+
+def _sdpa(q, k, v, mask):
+    # (B, L, H, D) in, (B, L, H * D) out; mask: bool (B, 1, 1, S), True = attend
+    q, k, v = (t.transpose(1, 2) for t in (q, k, v))
+    out = F.scaled_dot_product_attention(
+        q, k.to(q.dtype), v.to(q.dtype), attn_mask=mask
+    )
+    return out.transpose(1, 2).flatten(2)
+
+
+def _modulate(h, values, span):
+    """h * value per token: values = (noisy,) for all tokens, or (noisy, clean)
+    with only tokens [span[0]:span[1]] noisy."""
+    if span is None:
+        return h * values[0][:, None]
+    noisy, clean = values
+    s, e = span
+    return torch.cat(
+        (
+            h[:, :s] * clean[:, None],
+            h[:, s:e] * noisy[:, None],
+            h[:, e:] * clean[:, None],
+        ),
+        1,
+    )
+
+
+class _JointAttention(nn.Module):
+    def __init__(self, dim, heads, eps):
+        super().__init__()
+        self.heads = heads
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        self.q_norm = _RMSNorm(dim // heads, eps)
+        self.k_norm = _RMSNorm(dim // heads, eps)
+
+    def forward(self, x, rope, mask):
+        q, k, v = self.qkv(x).unflatten(-1, (3, self.heads, -1)).unbind(2)
+        q = _apply_rope(self.q_norm(q), rope)
+        k = _apply_rope(self.k_norm(k), rope)
+        return self.out(_sdpa(q, k, v, mask))
+
+
+class _Attention(nn.Module):
+    # no RoPE: the cross-attention to the text, and the Perceiver's self-attention
+    def __init__(self, dim, heads, kv_dim, eps):
+        super().__init__()
+        self.heads = heads
+        self.to_q = nn.Linear(dim, dim, bias=False)
+        self.to_kv = nn.Linear(kv_dim, 2 * dim, bias=False)
+        self.to_out = nn.Linear(dim, dim, bias=False)
+        self.norm_q = _RMSNorm(dim // heads, eps)
+        self.norm_k = _RMSNorm(dim // heads, eps)
+
+    def forward(self, x, context, mask):
+        q = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1)))
+        k, v = self.to_kv(context).unflatten(-1, (2, self.heads, -1)).unbind(2)
+        return self.to_out(_sdpa(q, self.norm_k(k), v, mask))
+
+
+class _FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.w13 = nn.Linear(dim, 2 * hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x):
+        gate, up = self.w13(x).chunk(2, -1)
+        return self.w2(F.silu(gate) * up)
+
+
+class _Block(nn.Module):
+    """Z-Image's block; an injection block also cross-attends to its bundle."""
+
+    def __init__(self, p: Fibo2Params, injection: bool):
+        super().__init__()
+        dim, eps = p.dim, p.norm_eps
+        self.attention = _JointAttention(dim, p.n_heads, eps)
+        self.feed_forward = _FeedForward(dim, p.ffn_dim)
+        self.attention_norm1 = _RMSNorm(dim, eps)
+        self.ffn_norm1 = _RMSNorm(dim, eps)
+        self.attention_norm2 = _RMSNorm(dim, eps)
+        self.ffn_norm2 = _RMSNorm(dim, eps)
+        self.adaLN_modulation = nn.Sequential(nn.Linear(ADALN_DIM, 4 * dim, bias=True))
+        self.cross_attention = None
+        if injection:
+            self.cross_attn_norm1 = _RMSNorm(dim, eps)
+            self.cross_attention = _Attention(dim, p.n_heads, p.cap_feat_dim, eps)
+            self.cross_attn_norm2 = _RMSNorm(dim, eps)
+            self.cross_attn_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, rope, adaln, span, mask=None, context=None, context_mask=None):
+        # adaln: (B, 256) for all tokens, or (2B, 256) = [noisy; clean] with span
+        mod = self.adaLN_modulation(adaln)
+        parts = (mod,) if span is None else mod.chunk(2, 0)
+        chunks = [m.chunk(4, -1) for m in parts]
+        scale_msa = [1 + c[0] for c in chunks]
+        gate_msa = [c[1].tanh() for c in chunks]
+        scale_mlp = [1 + c[2] for c in chunks]
+        gate_mlp = [c[3].tanh() for c in chunks]
+
+        y = self.attention(
+            _modulate(self.attention_norm1(x), scale_msa, span), rope, mask
+        )
+        x = x + _modulate(self.attention_norm2(y), gate_msa, span)
+        if self.cross_attention is not None:
+            y = self.cross_attention(self.cross_attn_norm1(x), context, context_mask)
+            x = x + self.cross_attn_gate.tanh() * self.cross_attn_norm2(y)
+        y = self.feed_forward(_modulate(self.ffn_norm1(x), scale_mlp, span))
+        return x + _modulate(self.ffn_norm2(y), gate_mlp, span)
+
+
+class _PerceiverLayer(nn.Module):
+    def __init__(self, p: Fibo2Params):
+        super().__init__()
+        dim, eps = p.dim, p.norm_eps
+        self.cross_attn_norm = _RMSNorm(dim, eps)
+        self.cross_attention = _Attention(dim, p.n_heads, p.cap_feat_dim, eps)
+        self.self_attn_norm = _RMSNorm(dim, eps)
+        self.self_attention = _Attention(dim, p.n_heads, dim, eps)
+        self.ffn_norm = _RMSNorm(dim, eps)
+        self.ffn = _FeedForward(dim, p.ffn_dim)
+
+    def forward(self, x, context, context_mask, gist_mask):
+        x = x + self.cross_attention(self.cross_attn_norm(x), context, context_mask)
+        y = self.self_attn_norm(x)
+        x = x + self.self_attention(y, y, gist_mask)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class _Perceiver(nn.Module):
+    def __init__(self, p: Fibo2Params):
+        super().__init__()
+        self.latent_queries = nn.Parameter(torch.zeros(p.max_gist_tokens, p.dim))
+        self.layers = nn.ModuleList(
+            [_PerceiverLayer(p) for _ in range(p.perceiver_layers)]
+        )
+        self.norm_out = _RMSNorm(p.dim, p.norm_eps)
+
+
+class _TimestepEmbedder(nn.Module):
+    def __init__(self, freq_dim=256, mid_dim=1024, out_dim=ADALN_DIM):
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_dim, mid_dim, bias=True),
+            nn.SiLU(),
+            nn.Linear(mid_dim, out_dim, bias=True),
+        )
+
+    def forward(self, t):
+        half = self.freq_dim // 2
+        freqs = torch.exp(
+            -math.log(10000)
+            * torch.arange(half, dtype=torch.float32, device=t.device)
+            / half
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat((args.cos(), args.sin()), -1)
+        return self.mlp(embedding.to(self.mlp[0].weight.dtype))
+
+
+class _FinalLayer(nn.Module):
+    def __init__(self, dim, out_channels):
+        super().__init__()
+        self.linear = nn.Linear(dim, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(ADALN_DIM, dim, bias=True)
+        )
+
+    def forward(self, x, adaln):
+        scale = 1 + self.adaLN_modulation(adaln)
+        x = F.layer_norm(x, (x.shape[-1],), eps=1e-6)
+        return self.linear(x * scale[:, None])
+
+
+class Fibo2Transformer(nn.Module, OstrisModelMixin):
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # unquantized tensors only (the block linears may be int8)
+        max_gist, dim = state_dict["perceiver.latent_queries"].shape
+        n_layers = 1 + max(
+            int(k.split(".")[1]) for k in state_dict if k.startswith("layers.")
+        )
+        return Fibo2Params(
+            dim=dim,
+            n_heads=dim // state_dict["layers.0.attention.q_norm.weight"].shape[0],
+            n_layers=n_layers,
+            n_refiner_layers=1
+            + max(
+                int(k.split(".")[1])
+                for k in state_dict
+                if k.startswith("noise_refiner.")
+            ),
+            injection_layers=tuple(
+                i
+                for i in range(n_layers)
+                if f"layers.{i}.cross_attn_gate" in state_dict
+            ),
+            # int8 convrot keeps the (out, in) weight shape
+            cap_feat_dim=state_dict[
+                "perceiver.layers.0.cross_attention.to_kv.weight"
+            ].shape[1],
+            perceiver_layers=1
+            + max(
+                int(k.split(".")[2])
+                for k in state_dict
+                if k.startswith("perceiver.layers.")
+            ),
+            max_gist_tokens=max_gist,
+            ffn_dim=state_dict["layers.0.feed_forward.w13.weight"].shape[0] // 2,
+            in_channels=state_dict["x_embedder.weight"].shape[1],
+        )
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["noise_refiner", "perceiver.layers", "layers"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # what the Comfy int8 convrot file keeps in bf16
+        return ["x_embedder", "t_embedder*", "final_layer*", "*adaLN_modulation*"]
+
+    def __init__(self, params: Optional[Fibo2Params] = None):
+        super().__init__()
+        p = params or Fibo2Params()
+        if p.dim // p.n_heads != 2 * sum(d // 2 for d in p.axes_dims):
+            raise ValueError("rope axes must cover the head dim")
+        self.params = p
+        self.x_embedder = nn.Linear(p.in_channels, p.dim, bias=True)
+        self.t_embedder = _TimestepEmbedder()
+        self.noise_refiner = nn.ModuleList(
+            [_Block(p, False) for _ in range(p.n_refiner_layers)]
+        )
+        self.perceiver = _Perceiver(p)
+        self.layers = nn.ModuleList(
+            [_Block(p, i in p.injection_layers) for i in range(p.n_layers)]
+        )
+        self.final_layer = _FinalLayer(p.dim, p.in_channels)
+        self.gradient_checkpointing = False
+
+    @property
+    def device(self):
+        return self.x_embedder.weight.device
+
+    @property
+    def dtype(self):
+        # never quantized
+        return self.x_embedder.weight.dtype
+
+    def enable_gradient_checkpointing(self, enable: bool = True):
+        self.gradient_checkpointing = enable
+
+    def disable_gradient_checkpointing(self):
+        self.gradient_checkpointing = False
+
+    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 _rope(self, ids: torch.Tensor):
+        """(L, 3) integer positions -> (cos, sin), each (1, L, 1, head_dim / 2)
+        fp32; phases in fp64 like diffusers' precomputed tables."""
+        angles = []
+        for i, d in enumerate(self.params.axes_dims):
+            freqs = 1.0 / (
+                self.params.rope_theta
+                ** (torch.arange(0, d, 2, dtype=torch.float64, device=ids.device) / d)
+            )
+            angles.append(torch.outer(ids[:, i].double(), freqs).float())
+        angles = torch.cat(angles, -1)[None, :, None]
+        return angles.cos(), angles.sin()
+
+    @staticmethod
+    def _grid(plane, height, width, device):
+        h, w = torch.meshgrid(
+            torch.arange(height, device=device),
+            torch.arange(width, device=device),
+            indexing="ij",
+        )
+        return torch.stack((torch.full_like(h, plane), h, w), -1).flatten(0, 1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # (B, 128, H, W) normalized FLUX.2 latents
+        timestep: torch.Tensor,  # 0..1000, 1000 = pure noise
+        encoder_hidden_states: torch.Tensor,  # (B, T, 6 * cap_feat_dim)
+        encoder_attention_mask: Optional[torch.Tensor] = None,  # (B, T) bool
+        ref_latents: Optional[List[torch.Tensor]] = None,  # each (B, 128, h, w)
+    ):
+        p = self.params
+        batch, channels, height, width = hidden_states.shape
+        device, dtype = hidden_states.device, self.dtype
+        num_image = height * width
+
+        t = (
+            torch.as_tensor(timestep, device=device, dtype=torch.float32)
+            .reshape(-1)
+            .expand(batch)
+        )
+        adaln_noisy = self.t_embedder(t)
+        adaln_both = torch.cat((adaln_noisy, self.t_embedder(torch.zeros_like(t))))
+
+        # the image on rope plane max_gist + 1 whatever the gist count, reference
+        # k on the k-th plane after it
+        plane = p.max_gist_tokens + 1
+        tokens = [hidden_states.flatten(2).transpose(1, 2)]
+        ids = [self._grid(plane, height, width, device)]
+        for k, ref in enumerate(ref_latents or []):
+            tokens.append(ref.flatten(2).transpose(1, 2))
+            ids.append(self._grid(plane + 1 + k, ref.shape[-2], ref.shape[-1], device))
+        x = self.x_embedder(torch.cat(tokens, 1).to(dtype))
+        image_rope = self._rope(torch.cat(ids))
+        span = (0, num_image) if ref_latents else None
+        for block in self.noise_refiner:
+            x = self._ckpt(
+                block, x, image_rope, adaln_both if span else adaln_noisy, span
+            )
+
+        context = encoder_hidden_states.to(dtype)
+        mask = encoder_attention_mask
+        if mask is not None:
+            mask = mask.to(device=device, dtype=torch.bool)
+            lengths = mask.sum(1).tolist()
+            context_mask = (
+                None
+                if all(n == mask.shape[1] for n in lengths)
+                else mask[:, None, None]
+            )
+        else:
+            lengths, context_mask = [context.shape[1]] * batch, None
+        bundles = context.chunk(1 + len(p.injection_layers), -1)
+
+        counts = gist_counts(lengths, p)
+        num_gist = max(counts)
+        gist_mask = None
+        if min(counts) < num_gist:
+            gist_mask = (
+                torch.arange(num_gist, device=device)
+                < torch.tensor(counts, device=device)[:, None]
+            )
+        gist = (
+            self.perceiver.latent_queries[:num_gist]
+            .to(dtype)
+            .unsqueeze(0)
+            .expand(batch, -1, -1)
+        )
+        for layer in self.perceiver.layers:
+            gist = self._ckpt(
+                layer,
+                gist,
+                bundles[0],
+                context_mask,
+                None if gist_mask is None else gist_mask[:, None, None],
+            )
+        gist = self.perceiver.norm_out(gist)
+        gist_ids = torch.zeros(num_gist, 3, dtype=torch.long, device=device)
+        gist_ids[:, 0] = torch.arange(1, num_gist + 1, device=device)
+        gist_rope = self._rope(gist_ids)
+
+        # gist first as trained; with padded gists last so the padding ends every row
+        if gist_mask is None:
+            x = torch.cat((gist, x), 1)
+            rope = tuple(torch.cat(r, 1) for r in zip(gist_rope, image_rope))
+            span, attn_mask = (num_gist, num_gist + num_image), None
+        else:
+            attn_mask = torch.cat(
+                (
+                    torch.ones(batch, x.shape[1], dtype=torch.bool, device=device),
+                    gist_mask,
+                ),
+                1,
+            )[:, None, None]
+            x = torch.cat((x, gist), 1)
+            rope = tuple(torch.cat(r, 1) for r in zip(image_rope, gist_rope))
+            span = (0, num_image)
+
+        text = {
+            layer_id: bundles[i + 1] for i, layer_id in enumerate(p.injection_layers)
+        }
+        for index, block in enumerate(self.layers):
+            x = self._ckpt(
+                block,
+                x,
+                rope,
+                adaln_both,
+                span,
+                attn_mask,
+                text.get(index),
+                context_mask,
+            )
+
+        x = self.final_layer(x[:, span[0] : span[1]], adaln_noisy)
+        return x.transpose(1, 2).reshape(batch, channels, height, width)

extensions_built_in/diffusion_models/ui.tsx

@@ -2177,6 +2177,52 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       "model.model_kwargs.kv_cache",
     ],
   },
+  {
+    name: "fibo2",
+    label: "Fibo 2.0 Dev",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      // the Comfy-Org repack (vendor repo is gated); a comfy-layout
+      // .safetensors path (fine-tune) loads as named
+      "config.process[0].model.name_or_path": [
+        "Comfy-Org/Fibo-2.0",
+        defaultNameOrPath,
+      ],
+      // convrot8 attaches the repack's int8 convrot DiT / Qwen3-VL-4B as-is
+      "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],
+      // vendor defaults: CFG 5 against the null caption, 30 steps
+      "config.process[0].sample.guidance_scale": [5, 4],
+      "config.process[0].sample.sample_steps": [30, 25],
+      "config.process[0].train.timestep_type": ["weighted", "sigmoid"],
+      "config.process[0].train.unload_text_encoder": [false, false],
+      "config.process[0].model.model_kwargs": [
+        {
+          match_target_res: true,
+        },
+        {},
+      ],
+    },
+    disableSections: ["network.conv", "train.unload_text_encoder"],
+    // one model: it edits when the dataset has control paths, and is plain
+    // text to image when it does not
+    additionalSections: [
+      "datasets.multi_control_paths",
+      "sample.multi_ctrl_imgs",
+      "model.qie.match_target_res",
+      "model.layer_offloading",
+      "model.low_vram",
+    ],
+    hasMultiLinePrompts: true,
+  },
   {
     name: "flux2_klein_4b",
     label: "FLUX.2-klein-base-4B",

extensions_built_in/inference_engine/latent_preview.py

@@ -53,6 +53,7 @@ ARCH_LATENT_FORMAT = {
     "prx_pixel": None,  # pixel space
     "flux3": "Flux3Image",
     "moss_image": "Flux2",  # FLUX.2 encoder, same packing
+    "fibo2": "Flux2",  # FLUX.2 encoder, same packing
     "iris3b": None,  # pixel space
 }
 

testing/test_model_loading.py

@@ -278,6 +278,29 @@ MODEL_TESTS = {
         },
         "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
     },
+    "fibo2": {
+        "model": {
+            "name_or_path": "Comfy-Org/Fibo-2.0",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        # vendor defaults are 30 steps, CFG 5
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 5.0},
+    },
+    "fibo2:edit": {
+        # same arch: a control image makes it an edit
+        "model": {
+            "name_or_path": "Comfy-Org/Fibo-2.0",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 5.0},
+        "needs_control_image": True,
+    },
     # ---- legacy monolith archs (components adopted into v2 on load) ----
     "sd1": {
         "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"},

toolkit/models/registry.py

@@ -60,6 +60,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "moss_image:base": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}},
     "moss_image:edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True, "uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}, "needs_control_image": True},
     "moss_image:base_edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True, "uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}, "needs_control_image": True},
+    "fibo2": {"modality": "image", "model": {"name_or_path": "Comfy-Org/Fibo-2.0", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 5.0}},
     "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}},

version.py

@@ -1 +1 @@
-VERSION = "0.13.36"
+VERSION = "0.13.37"