diff --git a/README.md b/README.md index 3862f50..c0825df 100644 --- a/README.md +++ b/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) diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 7b4ba8f..5dfad3a 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/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, ] diff --git a/extensions_built_in/diffusion_models/fibo2/__init__.py b/extensions_built_in/diffusion_models/fibo2/__init__.py new file mode 100644 index 0000000..b27dc85 --- /dev/null +++ b/extensions_built_in/diffusion_models/fibo2/__init__.py @@ -0,0 +1,3 @@ +from .fibo2 import Fibo2Model + +__all__ = ["Fibo2Model"] diff --git a/extensions_built_in/diffusion_models/fibo2/fibo2.py b/extensions_built_in/diffusion_models/fibo2/fibo2.py new file mode 100644 index 0000000..79ffb5b --- /dev/null +++ b/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() diff --git a/extensions_built_in/diffusion_models/fibo2/src/__init__.py b/extensions_built_in/diffusion_models/fibo2/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/fibo2/src/checkpoints.py b/extensions_built_in/diffusion_models/fibo2/src/checkpoints.py new file mode 100644 index 0000000..213eb18 --- /dev/null +++ b/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 + ) diff --git a/extensions_built_in/diffusion_models/fibo2/src/pipeline.py b/extensions_built_in/diffusion_models/fibo2/src/pipeline.py new file mode 100644 index 0000000..3da38bf --- /dev/null +++ b/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] diff --git a/extensions_built_in/diffusion_models/fibo2/src/references.py b/extensions_built_in/diffusion_models/fibo2/src/references.py new file mode 100644 index 0000000..ad253f4 --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/fibo2/src/text_encoder.py b/extensions_built_in/diffusion_models/fibo2/src/text_encoder.py new file mode 100644 index 0000000..37b15ac --- /dev/null +++ b/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 diff --git a/extensions_built_in/diffusion_models/fibo2/src/transformer.py b/extensions_built_in/diffusion_models/fibo2/src/transformer.py new file mode 100644 index 0000000..5b35682 --- /dev/null +++ b/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) diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 53948b5..85dbf43 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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", diff --git a/extensions_built_in/inference_engine/latent_preview.py b/extensions_built_in/inference_engine/latent_preview.py index 13209af..f98fe63 100644 --- a/extensions_built_in/inference_engine/latent_preview.py +++ b/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 } diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index d46ea60..6d24ef2 100644 --- a/testing/test_model_loading.py +++ b/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"}, diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index a243626..1dfeb8c 100644 --- a/toolkit/models/registry.py +++ b/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}}, diff --git a/version.py b/version.py index e7d4575..92152c0 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.36" +VERSION = "0.13.37"