diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 660f0e2..1075015 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -22,6 +22,7 @@ from .krea2 import Krea2Model from .boogu_image import BooguImageModel, BooguImageEditModel from .mageflow import MageFlowModel, MageFlowEditModel from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel +from .flux3 import Flux3Model AI_TOOLKIT_MODELS = [ # put a list of models here @@ -63,4 +64,5 @@ AI_TOOLKIT_MODELS = [ MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel, + Flux3Model, ] diff --git a/extensions_built_in/diffusion_models/flux3/__init__.py b/extensions_built_in/diffusion_models/flux3/__init__.py new file mode 100644 index 0000000..48ead9a --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/__init__.py @@ -0,0 +1,3 @@ +from .flux3 import Flux3Model + +__all__ = ["Flux3Model"] diff --git a/extensions_built_in/diffusion_models/flux3/flux3.py b/extensions_built_in/diffusion_models/flux3/flux3.py new file mode 100644 index 0000000..effb320 --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/flux3.py @@ -0,0 +1,535 @@ +"""FLUX.3 (small, ~5.5B) for ai-toolkit. + +A joint flow-matching DiT (``src/transformer.py``) that trains and samples +t2i and i2i (reference images as clean conditioning tokens). + + - image VAE: FLUX.2 packed AE, f16, 128 latent channels + - text: Qwen3-VL-4B, 8 hidden layers stacked along the channel dim -> + ``(B, L, 20480)``; the DiT's ``timesteps_ctx`` are zeros and there is no + pooled vector. Embeds are stored per-sample as 2D ``(L, 20480)`` tensors + in an ``AdvancedPromptEmbeds`` (padded to a multiple of 80 tokens -- the + DiT has no text mask, pad tokens are part of the trained context). + +``model.name_or_path`` is the model package (hub repo id or local dir): + + config.json model.safetensors ae.safetensors text_encoder/ tokenizer/ + +Only the DiT and the image VAE come from the package. The package's +text_encoder/ and tokenizer/ are byte-identical to Qwen/Qwen3-VL-4B-Instruct, +so the text encoder loads from the Comfy-Org single file under +MODELS_PATH/text_encoders/ (downloaded there once when missing) with config +and tokenizer from the public Qwen repo. + +Flow convention matches ai-toolkit exactly (t=1 noise -> t=0 clean, +target = noise - clean), so no time flips anywhere. +""" + +import json +import os +from typing import TYPE_CHECKING, List, Optional + +import torch +from PIL import Image +from safetensors.torch import load_file, save_file +from torchvision.transforms.functional import to_tensor + +from toolkit.accelerator import unwrap_model +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.metadata import get_meta_for_safetensors +from toolkit.models.base_model import BaseModel +from toolkit.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) + +from .src.image_vae import Flux3ImageVAE +from .src.pipeline import ( + DEFAULT_ALPHA, + Flux3Pipeline, + build_text_conditioning, + pack_i2i_image_condition, + pack_image_latents, +) +from .src.text_encoder import ( + DEFAULT_MAX_TOKENS, + DEFAULT_PAD_MULTIPLE, + QWEN3_VL_4B_REPO, + Flux3TextEncoder, + encode_flux3_prompt, +) +from .src.transformer import Flux3Transformer, Flux3TransformerParams + +if TYPE_CHECKING: + from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + + +# Static flow-matching timeshift used for the training timestep/noise +# schedule (the package's sampler alpha). +DEFAULT_TIMESTEP_SHIFT = DEFAULT_ALPHA + +scheduler_config = { + "num_train_timesteps": 1000, + "use_dynamic_shifting": False, +} + + +class Flux3Model(BaseModel): + arch = "flux3" + + def __init__( + self, + device, + model_config: ModelConfig, + dtype="bf16", + custom_pipeline=None, + noise_scheduler=None, + **kwargs, + ): + super().__init__( + device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs + ) + self.is_flow_matching = True + self.is_transformer = True + self.target_lora_modules = ["Flux3Transformer"] + self.supports_model_paths = True + self.use_old_lokr_format = False + + self.processor = None # Qwen3-VL AutoProcessor + self.max_text_length = int( + self.model_config.model_kwargs.get("max_text_length", DEFAULT_MAX_TOKENS) + ) + self.text_pad_multiple = int( + self.model_config.model_kwargs.get("text_pad_multiple", DEFAULT_PAD_MULTIPLE) + ) + self.package_config = None + self.latent_space_version = "flux3_v1" + + @staticmethod + def get_train_scheduler(): + return CustomFlowMatchEulerDiscreteScheduler( + shift=DEFAULT_TIMESTEP_SHIFT, **scheduler_config + ) + + def get_bucket_divisibility(self): + return 16 # f8 conv + 2x2 packing + + # ------------------------------------------------------------------ + # Package files + # ------------------------------------------------------------------ + def _package_file(self, filename: str, required: bool = True) -> Optional[str]: + """A package file: model_kwargs ``_path`` override, else the file + inside a local package dir, else downloaded from the hub repo.""" + override = self.model_config.model_kwargs.get( + filename.split(".")[0] + "_path", None + ) + if override is not None: + if not os.path.exists(override): + raise FileNotFoundError(f"FLUX.3 {filename} override not found: {override}") + return override + base = self.model_config.name_or_path + if os.path.isdir(base): + path = os.path.join(base, filename) + if os.path.exists(path): + return path + if required: + raise FileNotFoundError(f"FLUX.3 package file not found: {path}") + return None + if os.path.isfile(base) and filename == "model.safetensors": + return base # name_or_path points straight at the DiT file + from huggingface_hub import hf_hub_download + + try: + return hf_hub_download( + repo_id=base, filename=filename, token=os.getenv("HF_TOKEN", None) + ) + except Exception as e: + if required: + raise + self.print_and_status_update(f" - no {filename} in {base} ({type(e).__name__})") + return None + + def _load_package_config(self) -> Optional[dict]: + if self.package_config is None: + path = self._package_file("config.json", required=False) + if path is not None: + with open(path, "r") as f: + self.package_config = json.load(f) + return self.package_config + + def _transformer_params(self) -> Flux3TransformerParams: + cfg = self._load_package_config() + if cfg is None: + return Flux3TransformerParams() + d = cfg["dit"] + in_channels = {m: c for m, c in d["in_channels"].items() if m != "text"} + return Flux3TransformerParams( + in_channels=in_channels, + sequence={f"x_{m}": m for m in in_channels}, + vec_in_dim=d.get("vec_in_dim", None), + context_in_dim=d["in_channels"]["text"], + hidden_size=d["hidden"], + num_heads=d["heads"], + depth=d["depth"], + depth_single_blocks=d["depth_single"], + axes_dim=list(d["axes_dim"]), + theta=d["theta"], + mlp_ratio=d["mlp_ratio"], + out_modes=list(d["out_modes"]), + ) + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + def _load_transformer(self) -> Flux3Transformer: + dtype = self.torch_dtype + self.print_and_status_update("Loading transformer (Flux3Transformer)") + dit_path = self._package_file("model.safetensors") + params = self._transformer_params() + cfg = self._load_package_config() + if cfg is not None and cfg.get("quantization", None): + raise ValueError( + "FLUX.3 package variants with pre-quantized weights " + f"({cfg['quantization']}) are not supported; use the bf16 package " + "and let the toolkit quantize." + ) + + state_dict = load_file(dit_path) + state_dict = { + k: (v.to(dtype) if v.is_floating_point() else v) + for k, v in state_dict.items() + } + with torch.device("meta"): + transformer = Flux3Transformer(params) + needed = set(transformer.state_dict().keys()) + missing = sorted(needed - set(state_dict.keys())) + if missing: + raise ValueError(f"FLUX.3 checkpoint missing required keys: {missing[:8]}") + result = transformer.load_state_dict(state_dict, strict=False, assign=True) + if result.unexpected_keys: + self.print_and_status_update( + f" - ignoring {len(result.unexpected_keys)} unexpected checkpoint " + f"keys (e.g. {result.unexpected_keys[:3]})" + ) + del state_dict + flush() + + attn_backend = self.model_config.model_kwargs.get("attention_backend", None) + if attn_backend is not None: + transformer.set_attention_backend(attn_backend) + return transformer + + def _load_text_encoder(self): + te_path = self.model_config.te_name_or_path + if te_path is None: + te_path = self.model_config.model_kwargs.get( + "text_encoder_path", QWEN3_VL_4B_REPO + ) + self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}") + te_kwargs = self.component_load_kwargs("te") + text_encoder = Flux3TextEncoder.load_model( + te_path, + dtype=te_kwargs["dtype"], + subfolder="", + qtype=te_kwargs["qtype"], # comfy candidate ranking only + quantize_on_load=False, + ) + processor_source = te_path if os.path.isdir(te_path) else None + processor = Flux3TextEncoder.load_processor(processor_source) + processor.tokenizer.padding_side = "right" + # text-only conditioning -- drop the vision tower + text_encoder.drop_vision_tower() + text_encoder.eval() + text_encoder.requires_grad_(False) + flush() + return processor, text_encoder + + def _load_vae(self) -> Flux3ImageVAE: + self.print_and_status_update("Loading image VAE") + vae = Flux3ImageVAE() + vae.load_state_dict(load_file(self._package_file("ae.safetensors"))) + vae.to(self.vae_torch_dtype).eval().requires_grad_(False) + return vae + + def load_model(self): + self.print_and_status_update("Loading FLUX.3 model") + + transformer = self._load_transformer() + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) + flush() + + processor, text_encoder = self._load_text_encoder() + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) + flush() + + vae = self._load_vae() + vae.to(self.vae_device_torch) + + self.noise_scheduler = Flux3Model.get_train_scheduler() + + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = processor.tokenizer + self.processor = processor + self.model = transformer + self.pipeline = Flux3Pipeline(self) + self.print_and_status_update("Model Loaded") + + # ------------------------------------------------------------------ + # Generation (training previews) + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return Flux3Pipeline(self) + + def generate_single_image( + self, + pipeline: Flux3Pipeline, + gen_config: GenerateImageConfig, + conditional_embeds: AdvancedPromptEmbeds, + unconditional_embeds: AdvancedPromptEmbeds, + generator: torch.Generator, + extra: dict, + ): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + sc = self.get_bucket_divisibility() + gen_config.width = int(gen_config.width // sc * sc) + gen_config.height = int(gen_config.height // sc * sc) + + # cond and uncond were encoded separately, so their pad-to-80 windows + # can differ. The DiT attends to real pad-token states, not zeros, so + # re-encode the shorter prompt at the longer window when the TE is up. + if unconditional_embeds is not None and self.text_encoder is not None: + cond_len = conditional_embeds.text_embeds[0].shape[0] + uncond_len = unconditional_embeds.text_embeds[0].shape[0] + if cond_len != uncond_len: + target = max(cond_len, uncond_len) + if cond_len < target: + conditional_embeds = self._encode_at_length(gen_config.prompt, target) + else: + unconditional_embeds = self._encode_at_length( + gen_config.negative_prompt or "", target + ) + + ctrl_img = None + if gen_config.ctrl_img is not None: + img = Image.open(gen_config.ctrl_img).convert("RGB") + img = img.resize((gen_config.width, gen_config.height), Image.LANCZOS) + ctrl_img = to_tensor(img).unsqueeze(0) + + alpha = self.model_config.model_kwargs.get("sample_alpha", None) + alpha = float(alpha) if alpha is not None else None + + result = pipeline( + conditional_embeds=conditional_embeds, + unconditional_embeds=unconditional_embeds, + height=gen_config.height, + width=gen_config.width, + num_inference_steps=gen_config.num_inference_steps, + guidance_scale=gen_config.guidance_scale, + latents=gen_config.latents, + generator=generator, + ctrl_img=ctrl_img, + alpha=alpha, + ) + return result[0] + + def _encode_at_length(self, prompt: str, length: int) -> AdvancedPromptEmbeds: + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + ctx = encode_flux3_prompt( + self.text_encoder, + self.processor, + [prompt], + max_length=length, + dtype=self.torch_dtype, + pad_multiple=None, + ) + return AdvancedPromptEmbeds(text_embeds=[ctx[0]]) + + # ------------------------------------------------------------------ + # Training hooks + # ------------------------------------------------------------------ + def get_noise_prediction( + self, + latent_model_input: torch.Tensor, # (B, 128, h, w) + timestep: torch.Tensor, # 0..1000 scale + text_embeddings: AdvancedPromptEmbeds, + batch: "DataLoaderBatchDTO" = None, + **kwargs, + ): + device = self.device_torch + dtype = self.torch_dtype + if self.model.device == torch.device("cpu"): + self.model.to(device) + + batch_size = latent_model_input.shape[0] + + with torch.no_grad(): + # toolkit timestep (0..1000, 1000 = pure noise) -> flow time in [0, 1] + t = timestep.to(device, dtype=torch.float32) / 1000.0 + if t.dim() == 0: + t = t.unsqueeze(0) + if t.shape[0] != batch_size: + t = t.expand(batch_size) + + ctx, ctx_ids, ts_ctx = build_text_conditioning( + text_embeddings.text_embeds, device, dtype + ) + + latent_model_input = latent_model_input.to(device, dtype) + _, _, h_lat, w_lat = latent_model_input.shape + tokens, ids = pack_image_latents(latent_model_input) + model_kwargs = { + "x_image": tokens.to(dtype), + "x_image_ids": ids, + "x_image_timesteps": t.to(dtype)[:, None] + * torch.ones(tokens.shape[:2], device=device, dtype=dtype), + } + + # i2i: reference image(s) as clean t=0 tokens at l = -1, -2, ... + control_list = None + if batch is not None: + if batch.control_tensor_list is not None: + control_list = batch.control_tensor_list + elif batch.control_tensor is not None: + control_list = [[batch.control_tensor[b]] for b in range(batch_size)] + if control_list is not None: + ref_latents = [] + for controls in control_list: + sample_refs = [] + for ctrl in controls: + if ctrl.ndim == 4: + ctrl = ctrl[0] + # control tensors arrive in [0, 1] + lat = self.encode_images( + ctrl.unsqueeze(0) * 2 - 1, device=device, dtype=dtype + ) + sample_refs.append(lat[0]) + ref_latents.append(sample_refs) + xc, xc_ids, xc_ts = pack_i2i_image_condition(ref_latents, device, dtype) + model_kwargs["x_image_cond"] = xc + model_kwargs["x_image_cond_ids"] = xc_ids + model_kwargs["x_image_cond_timesteps"] = xc_ts + + pred = self.transformer( + ctx=ctx, + ctx_ids=ctx_ids, + timesteps_ctx=ts_ctx, + **model_kwargs, + ) + out = pred["x_image"] + return out.reshape(batch_size, h_lat, w_lat, -1).permute(0, 3, 1, 2) + + def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds: + if isinstance(prompt, str): + prompt = [prompt] + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + + ctx = encode_flux3_prompt( + self.text_encoder, + self.processor, + prompt, + max_length=self.max_text_length, + dtype=self.torch_dtype, + pad_multiple=self.text_pad_multiple, + ) + # one 2D (L, 20480) tensor per batch item + return AdvancedPromptEmbeds(text_embeds=[ctx[i] for i in range(ctx.shape[0])]) + + def get_loss_target(self, *args, **kwargs): + noise = kwargs.get("noise") + batch = kwargs.get("batch") + return (noise - batch.latents).detach() + + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + # ------------------------------------------------------------------ + # VAE encode / decode + # ------------------------------------------------------------------ + @torch.no_grad() + def encode_images(self, image_list, device=None, dtype=None): + """Images (C, H, W) in [-1, 1] -> latents (B, 128, H/16, W/16). + + Accepts a list of per-item tensors or a stacked batch tensor. + """ + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + + vae_dtype = self.vae.dtype + batch = torch.stack([img.to(device, dtype=vae_dtype) for img in image_list]) + latents = self.vae.encode(batch) + return latents.to(device, dtype=dtype) + + @torch.no_grad() + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + images = self.vae.decode(latents.to(device, self.vae.dtype)) + return images.to(device, dtype) + + # ------------------------------------------------------------------ + # Saving / bookkeeping + # ------------------------------------------------------------------ + def save_model(self, output_path, meta, save_dtype): + from toolkit.util.quantize import dequantize_if_quantized + + if not output_path.endswith(".safetensors"): + output_path = output_path + ".safetensors" + transformer: Flux3Transformer = unwrap_model(self.model) + state_dict = transformer.state_dict() + save_dict = {} + for k, v in state_dict.items(): + save_dict[k] = ( + dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype) + ) + meta = get_meta_for_safetensors(meta, name=self.arch) + save_file(save_dict, output_path, metadata=meta) + + def get_base_model_version(self): + return "flux3" + + def get_transformer_block_names(self) -> Optional[List[str]]: + return Flux3Transformer.get_transformer_block_names() + + def get_quantization_exclude_modules(self): + # sensitive small modules kept in full precision (fnmatch patterns): + # emb_in*/txt_in - token input projections + # time_in* - timestep embedder; feeds every modulation + # *modulations* - per-stream adaLN modulations + # final_layer* - modulated output projections + return [ + "emb_in*", + "txt_in", + "time_in*", + "*modulations*", + "final_layer*", + ] + + def convert_lora_weights_before_save(self, state_dict): + return { + k.replace("transformer.", "diffusion_model."): v + for k, v in state_dict.items() + } + + def convert_lora_weights_before_load(self, state_dict): + return { + k.replace("diffusion_model.", "transformer."): v + for k, v in state_dict.items() + } diff --git a/extensions_built_in/diffusion_models/flux3/src/__init__.py b/extensions_built_in/diffusion_models/flux3/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/flux3/src/image_vae.py b/extensions_built_in/diffusion_models/flux3/src/image_vae.py new file mode 100644 index 0000000..16f98d4 --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/image_vae.py @@ -0,0 +1,361 @@ +"""FLUX.3 image VAE -- the FLUX.2 packed image autoencoder. + +f8 conv autoencoder with 32 z-channels whose mean latent is 2x2 +space-to-channel packed to ``(B, 128, H/16, W/16)`` and normalized by a +frozen BatchNorm's running stats. ``encode``/``decode`` operate on +``(B, 3, H, W)`` pixels in ``[-1, 1]``. +""" + +import math +from dataclasses import dataclass, field + +import torch +from einops import rearrange +from torch import Tensor, nn + + +@dataclass +class Flux3ImageVAEParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + out_ch: int = 3 + ch_mult: list[int] = field(default_factory=lambda: [1, 2, 4, 4]) + num_res_blocks: int = 2 + z_channels: int = 32 + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, h, w = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + self.conv1 = nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + self.norm2 = nn.GroupNorm( + num_groups=32, num_channels=out_channels, eps=1e-6, affine=True + ) + self.conv2 = nn.Conv2d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=2, padding=0 + ) + + def forward(self, x: Tensor): + pad = (0, 1, 0, 1) + x = nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, x: Tensor): + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + x = self.conv(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.quant_conv = torch.nn.Conv2d(2 * z_channels, 2 * z_channels, 1) + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.conv_in = nn.Conv2d( + in_channels, self.ch, kernel_size=3, stride=1, padding=1 + ) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + curr_res = curr_res // 2 + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + self.norm_out = nn.GroupNorm( + num_groups=32, num_channels=block_in, eps=1e-6, affine=True + ) + self.conv_out = nn.Conv2d( + block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, x: Tensor) -> Tensor: + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + h = self.quant_conv(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.post_quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1) + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.ffactor = 2 ** (self.num_resolutions - 1) + + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + + self.conv_in = nn.Conv2d( + z_channels, block_in, kernel_size=3, stride=1, padding=1 + ) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + curr_res = curr_res * 2 + self.up.insert(0, up) + + self.norm_out = nn.GroupNorm( + num_groups=32, num_channels=block_in, eps=1e-6, affine=True + ) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + z = self.post_quant_conv(z) + + upscale_dtype = next(self.up.parameters()).dtype + + h = self.conv_in(z) + + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + + h = h.to(upscale_dtype) + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class Flux3ImageVAE(nn.Module): + def __init__(self, params: Flux3ImageVAEParams | None = None): + super().__init__() + if params is None: + params = Flux3ImageVAEParams() + self.params = params + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + + self.bn_eps = 1e-4 + self.bn_momentum = 0.1 + self.ps = [2, 2] + self.bn = torch.nn.BatchNorm2d( + math.prod(self.ps) * params.z_channels, + eps=self.bn_eps, + momentum=self.bn_momentum, + affine=False, + track_running_stats=True, + ) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def normalize(self, z: Tensor) -> Tensor: + self.bn.eval() + return self.bn(z) + + def inv_normalize(self, z: Tensor) -> Tensor: + self.bn.eval() + s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn_eps) + m = self.bn.running_mean.view(1, -1, 1, 1) + return z * s + m + + def encode(self, x: Tensor) -> Tensor: + moments = self.encoder(x) + mean = torch.chunk(moments, 2, dim=1)[0] + + z = rearrange( + mean, + "... c (i pi) (j pj) -> ... (c pi pj) i j", + pi=self.ps[0], + pj=self.ps[1], + ) + z = self.normalize(z) + return z + + def decode(self, z: Tensor) -> Tensor: + z = self.inv_normalize(z) + z = rearrange( + z, + "... (c pi pj) i j -> ... c (i pi) (j pj)", + pi=self.ps[0], + pj=self.ps[1], + ) + dec = self.decoder(z) + return dec diff --git a/extensions_built_in/diffusion_models/flux3/src/pipeline.py b/extensions_built_in/diffusion_models/flux3/src/pipeline.py new file mode 100644 index 0000000..ab935dc --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/pipeline.py @@ -0,0 +1,223 @@ +"""FLUX.3 sampling pipeline + shared latent/conditioning packers. + +Covers t2i and i2i (reference images). All conditioning arrives as prompt +*embeds* (never raw text) plus optional control-image tensors. + +Flow convention: time runs t=1 (pure noise) -> t=0 (clean); the DiT +predicts the velocity ``noise - clean`` -- identical to ai-toolkit's +convention. The inference schedule is a timeshifted uniform grid +(``timeshift(alpha, t)``); the package default alpha is 6.93. +""" + +from typing import List, Optional + +import torch +import torch.nn.functional as F +from PIL import Image +from diffusers.utils.torch_utils import randn_tensor + +from .positional import batched_prc_img, prc_img, prc_txt + +IMAGE_ROPE_SCALE = 0.5 +I2I_COND_L = -1 +DEFAULT_ALPHA = 6.93 +DEFAULT_GUIDANCE = 3.5 + + +def timeshift(alpha: float, t: torch.Tensor) -> torch.Tensor: + return alpha * t / (1.0 + (alpha - 1.0) * t) + + +def build_timesteps(n_steps: int, alpha: float) -> torch.Tensor: + return timeshift(alpha, torch.linspace(1.0, 0.0, n_steps + 1)) + + +# --------------------------------------------------------------------------- +# Text conditioning +# --------------------------------------------------------------------------- + + +def build_text_conditioning( + text_embeds_list: List[torch.Tensor], + device: torch.device, + dtype: torch.dtype, +): + """Per-sample (L, 20480) ctx tensors -> batched DiT text kwargs. + + Embeds are padded to a multiple of 80 tokens at encode time, so items + encoded separately (TE cache, cond/uncond pairs) can differ in length; + shorter ones are zero-padded to the batch's longest. Returns + ``(ctx (B, L, D), ctx_ids (B, L, 4), timesteps_ctx (B, L))``. + """ + max_len = max(e.shape[0] for e in text_embeds_list) + ctx = torch.stack( + [ + F.pad(e.to(device, dtype), (0, 0, 0, max_len - e.shape[0])) + for e in text_embeds_list + ] + ) + _, ctx_ids = prc_txt(ctx[0]) + ctx_ids = ctx_ids.unsqueeze(0).expand(ctx.shape[0], -1, -1).to(device) + timesteps_ctx = torch.zeros(ctx.shape[:2], device=device, dtype=dtype) + return ctx, ctx_ids, timesteps_ctx + + +# --------------------------------------------------------------------------- +# Token packing (training + sampling share these) +# --------------------------------------------------------------------------- + + +def pack_image_latents(latents: torch.Tensor): + """(B, 128, h, w) image latents -> ((B, h*w, 128) tokens, float ids with + the image RoPE scale applied to the h/w axes).""" + tokens, ids = batched_prc_img(latents) + ids = ids.to(dtype=torch.float32) + ids[..., 1:3] = ids[..., 1:3] * IMAGE_ROPE_SCALE + return tokens, ids + + +def pack_i2i_image_condition( + latents_per_sample: List[List[torch.Tensor]], + device: torch.device, + dtype: torch.dtype, +): + """Reference-image latents for i2i editing. + + ``latents_per_sample`` is one list of (128, h, w) latents per batch item + (every item must have the same reference geometry). Reference i placed at + ``t = 0, l = -1 - i``; conditioning tokens run at timestep 0. + Returns ``(tokens, ids, timesteps)``. + """ + all_tokens, all_ids = [], [] + for refs in latents_per_sample: + parts, id_parts = [], [] + for i_ref, latent in enumerate(refs): + x_ref, x_ref_ids = prc_img( + latent, + t_coord=torch.tensor([0], device=latent.device, dtype=torch.long), + l_coord=torch.tensor( + [I2I_COND_L - i_ref], device=latent.device, dtype=torch.long + ), + ) + parts.append(x_ref) + id_parts.append(x_ref_ids) + all_tokens.append(torch.cat(parts, dim=0)) + all_ids.append(torch.cat(id_parts, dim=0)) + tokens = torch.stack(all_tokens).to(device, dtype) + ids = torch.stack(all_ids).to(device=device, dtype=torch.float32) + ids[..., 1:3] = ids[..., 1:3] * IMAGE_ROPE_SCALE + timesteps = torch.zeros(tokens.shape[:2], device=device, dtype=dtype) + return tokens, ids, timesteps + + +# --------------------------------------------------------------------------- +# Sampling pipeline +# --------------------------------------------------------------------------- + + +class Flux3Pipeline: + """Lightweight sampler used by ai-toolkit's preview generation.""" + + def __init__(self, model): + # ``model`` is the Flux3Model so we can reuse its encode/decode + config. + self.model = model + + @property + def device(self): + return self.model.device_torch + + def to(self, *args, **kwargs): + return self + + def set_progress_bar_config(self, **kwargs): + pass + + @torch.no_grad() + def __call__( + self, + conditional_embeds, + unconditional_embeds, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = DEFAULT_GUIDANCE, + latents: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + ctrl_img: Optional[torch.Tensor] = None, # (1, 3, H, W) in [0, 1] + alpha: Optional[float] = None, + **kwargs, + ): + model = self.model + device = model.device_torch + dtype = model.torch_dtype + transformer = model.transformer + + do_cfg = unconditional_embeds is not None and guidance_scale > 1.0 + + if alpha is None: + alpha = DEFAULT_ALPHA + timesteps = build_timesteps(num_inference_steps, alpha=alpha).to(device) + + # text conditioning, CFG-doubled as [uncond, cond] + cond_list = conditional_embeds.text_embeds + if do_cfg: + embeds_list = unconditional_embeds.text_embeds + cond_list + else: + embeds_list = cond_list + ctx, ctx_ids, ts_ctx = build_text_conditioning(embeds_list, device, dtype) + cfg_batch = ctx.shape[0] + + # initial noise tokens + h_lat, w_lat = height // 16, width // 16 + if latents is None: + latents = randn_tensor( + (1, 128, h_lat, w_lat), generator=generator, dtype=torch.float32 + ) + latents = latents.to(device) + x, ids = pack_image_latents(latents) + x = x.float() + ids = ids.repeat(cfg_batch, 1, 1).to(device) + + # reference-image conditioning + cond_kwargs = {} + if ctrl_img is not None: + ctrl = ctrl_img.to(device, dtype) * 2 - 1 + cond_latent = model.encode_images(list(ctrl), device=device, dtype=dtype) + xc, xc_ids, xc_ts = pack_i2i_image_condition( + [[cond_latent[0]]], device, dtype + ) + cond_kwargs = { + "x_image_cond": xc.repeat(cfg_batch, 1, 1), + "x_image_cond_ids": xc_ids.repeat(cfg_batch, 1, 1), + "x_image_cond_timesteps": xc_ts.repeat(cfg_batch, 1), + } + + # Euler integration of the flow ODE with CFG + for step_index, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + x_in = x.to(dtype).repeat(cfg_batch, 1, 1) + pred = transformer( + ctx=ctx, + ctx_ids=ctx_ids, + timesteps_ctx=ts_ctx, + x_image=x_in, + x_image_ids=ids, + x_image_timesteps=t_curr + * torch.ones(x_in.shape[:2], device=device, dtype=dtype), + **cond_kwargs, + ) + v = pred["x_image"].float() + if do_cfg: + v_uc, v_c = v.chunk(2) + v = v_uc + guidance_scale * (v_c - v_uc) + if getattr(model, "sample_step_hook", None) is not None: + # inference engine preview: x0 estimate + x0 = (x - t_curr * v).reshape(1, h_lat, w_lat, -1).permute(0, 3, 1, 2) + model._emit_sample_step(x0, step_index, len(timesteps) - 1) + x = x + (t_prev - t_curr) * v + + img_latent = x.reshape(1, h_lat, w_lat, -1).permute(0, 3, 1, 2) + images = model.decode_latents(img_latent, device=device, dtype=dtype) + images = images.float().clamp(-1.0, 1.0) + images = ((images + 1.0) * 127.5).round().to(torch.uint8) + images = images.permute(0, 2, 3, 1).cpu().numpy() + return [Image.fromarray(arr) for arr in images] diff --git a/extensions_built_in/diffusion_models/flux3/src/positional.py b/extensions_built_in/diffusion_models/flux3/src/positional.py new file mode 100644 index 0000000..f0f3b8f --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/positional.py @@ -0,0 +1,69 @@ +"""Position-id packers for the FLUX.3 joint sequence. + +Every token stream fed to the DiT carries a 4-axis position id ``(t, h, w, l)``. +Image tokens use ``h``/``w`` (``t = 0``); the ``l`` axis separates reference +images (i2i conditioning uses ``l = -1, -2, ...``) and carries the sequence +position for text tokens. +""" + +from typing import Callable + +import torch +from einops import rearrange +from torch import Tensor + + +def prc_img( + x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None +) -> tuple[Tensor, Tensor]: + """(C, H, W) latent -> ((H*W, C) tokens, (H*W, 4) ids).""" + _, h, w = x.shape + dev = x.device + coords = { + "t": torch.arange(1, device=dev) if t_coord is None else t_coord.to(dev), + "h": torch.arange(h, device=dev), + "w": torch.arange(w, device=dev), + "l": torch.arange(1, device=dev) if l_coord is None else l_coord.to(dev), + } + x_ids = torch.cartesian_prod(coords["t"], coords["h"], coords["w"], coords["l"]) + x = rearrange(x, "c h w -> (h w) c") + return x, x_ids + + +def prc_txt( + x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None +) -> tuple[Tensor, Tensor]: + """(L, D) text context -> ((L, D) tokens, (L, 4) ids) with l = arange(L).""" + assert l_coord is None, "l_coord not supported for txt" + length, _ = x.shape + dev = x.device + coords = { + "t": torch.arange(1, device=dev) if t_coord is None else t_coord.to(dev), + "h": torch.arange(1, device=dev), + "w": torch.arange(1, device=dev), + "l": torch.arange(length, device=dev), + } + x_ids = torch.cartesian_prod(coords["t"], coords["h"], coords["w"], coords["l"]) + return x, x_ids + + +def _batched(fn: Callable) -> Callable: + def wrapped( + x: Tensor, t_coord: Tensor | None = None, l_coord: Tensor | None = None + ) -> tuple[Tensor, Tensor]: + pairs = [ + fn( + x[i], + t_coord[i] if t_coord is not None else None, + l_coord[i] if l_coord is not None else None, + ) + for i in range(len(x)) + ] + xs, ids = zip(*pairs) + return torch.stack(xs), torch.stack(ids) + + return wrapped + + +batched_prc_img = _batched(prc_img) +batched_prc_txt = _batched(prc_txt) diff --git a/extensions_built_in/diffusion_models/flux3/src/text_encoder.py b/extensions_built_in/diffusion_models/flux3/src/text_encoder.py new file mode 100644 index 0000000..064a560 --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/text_encoder.py @@ -0,0 +1,123 @@ +"""FLUX.3 text conditioning -- Qwen3-VL-4B hidden-state stack. + +The DiT text context is the stack of a fixed set of Qwen3-VL hidden layers +along the channel dim: ``ctx = (B, L, len(SELECT_LAYERS) * 2560 = 20480)``. +Prompts go through the chat template and are right-padded to a multiple of +``pad_multiple`` tokens (the package's ``text_encoder.pad_multiple``); the +model was trained with the pad tokens in-context and the DiT has no text +attention mask. FLUX.3 feeds the DiT zero ``timesteps_ctx`` and no pooled +vector. +""" + +import math +from typing import List, Optional + +import torch +from einops import rearrange + +from toolkit.models.v2.text_encoders.qwen3_vl import Qwen3VLTextEncoder + +SELECT_LAYERS = [4, 8, 12, 16, 20, 24, 28, 32] +CONTEXT_DIM = 20480 # len(SELECT_LAYERS) * qwen hidden (2560) + +# package defaults (text_encoder section of config.json) +DEFAULT_PAD_MULTIPLE = 80 +DEFAULT_MAX_TOKENS = 8192 + +QWEN3_VL_4B_REPO = "Qwen/Qwen3-VL-4B-Instruct" + + +class Flux3TextEncoder(Qwen3VLTextEncoder): + """Qwen3-VL-4B-Instruct. The package's text_encoder/ and tokenizer/ are + byte-identical to the public Qwen repo, so config + tokenizer come from + there and the weights from the Comfy-Org single-file repack under + MODELS_PATH/text_encoders/.""" + + aitk_config_repo = QWEN3_VL_4B_REPO + aitk_tokenizer_repo = QWEN3_VL_4B_REPO + aitk_tokenizer_subfolder = "" + aitk_processor_repo = QWEN3_VL_4B_REPO + aitk_processor_subfolder = "" + + aitk_comfy_repo = "Comfy-Org/Qwen3-VL" + _COMFY_FILES = [ + "text_encoders/qwen3vl_4b_bf16.safetensors", + "text_encoders/qwen3vl_4b_fp8_scaled.safetensors", + ] + aitk_comfy_weight_names = {QWEN3_VL_4B_REPO: _COMFY_FILES} + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + # comfy files drop the tied lm_head; strict load wants the key + embed = state_dict.get("model.language_model.embed_tokens.weight") + if embed is not None and "lm_head.weight" not in state_dict: + state_dict["lm_head.weight"] = embed + return state_dict + + +def padded_token_length( + tokenizer, + texts: List[str], + pad_multiple: int, + max_tokens: int, +) -> int: + """Common padded length for a batch: the longest prompt rounded up to a + multiple of ``pad_multiple``, capped at ``max_tokens``.""" + longest = 0 + for text in texts: + real = tokenizer( + text, return_tensors="pt", truncation=True, max_length=max_tokens + )["input_ids"].shape[1] + longest = max(longest, real) + return min(math.ceil(longest / pad_multiple) * pad_multiple, max_tokens) + + +@torch.no_grad() +def encode_flux3_prompt( + text_encoder, + processor, + prompts: List[str], + max_length: int = DEFAULT_MAX_TOKENS, + dtype: torch.dtype = torch.bfloat16, + pad_multiple: Optional[int] = DEFAULT_PAD_MULTIPLE, +) -> torch.Tensor: + """Prompts -> text context ``(B, L, 20480)``. + + ``text_encoder`` is a ``Qwen3VLForConditionalGeneration``; ``processor`` the + matching ``AutoProcessor``. Prompts are wrapped in the chat template with a + generation prompt and right-padded (matching how the DiT was trained -- do + not strip padding). ``L`` is the batch's longest prompt rounded up to a + multiple of ``pad_multiple``, capped at ``max_length``; with + ``pad_multiple=None`` every prompt is padded to exactly ``max_length``. + """ + messages = [[{"role": "user", "content": p}] for p in prompts] + texts = [ + processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) + for msg in messages + ] + tokenizer = processor.tokenizer + if pad_multiple is not None: + max_length = padded_token_length(tokenizer, texts, pad_multiple, max_length) + input_ids_list, attn_list = [], [] + for text in texts: + toks = tokenizer( + text, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_length, + ) + input_ids_list.append(toks["input_ids"]) + attn_list.append(toks["attention_mask"]) + input_ids = torch.cat(input_ids_list, dim=0).to(text_encoder.device) + attention_mask = torch.cat(attn_list, dim=0).to(text_encoder.device) + + out = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + stacked = torch.stack([out.hidden_states[k] for k in SELECT_LAYERS], dim=1) + result = rearrange(stacked, "b c l d -> b l (c d)") + return result.to(dtype) diff --git a/extensions_built_in/diffusion_models/flux3/src/transformer.py b/extensions_built_in/diffusion_models/flux3/src/transformer.py new file mode 100644 index 0000000..3a6c640 --- /dev/null +++ b/extensions_built_in/diffusion_models/flux3/src/transformer.py @@ -0,0 +1,725 @@ +"""Flux3Transformer -- the FLUX.3 joint single-sequence DiT. + +Every content mode is a token stream: ``x_image`` is the denoised stream and +``x_image_cond`` carries clean (t=0) reference tokens. Each stream is a +``(B, S, C_mode)`` token sequence with 4-axis RoPE position ids ``(B, S, 4)`` +and per-token timesteps ``(B, S)``. Text context ``ctx`` is ``(B, L, 20480)`` +(stacked Qwen3-VL hidden layers); FLUX.3 has no pooled text vector. + +Each mode gets its own input projection, ``depth`` early self-attention +blocks and modulation tables; text gets the same. The active modes and text +are then concatenated into one sequence for ``depth_single_blocks`` joint +blocks, split back apart, and the output modes go through their own final +layer. A mode with no tokens runs a 2-token zero dummy through its early +blocks (keeps the graph uniform) and never enters the joint sequence. + +The module tree matches the shipped ``model.safetensors`` layout so the state +dict loads without any key renaming. + +Attention defaults to torch's native ``scaled_dot_product_attention`` so it +runs on any GPU; ``set_attention_backend`` can opt into the cuDNN-pinned +kernel (Blackwell) or flash-attn 3 when installed. +""" + +import math +from collections import defaultdict +from dataclasses import dataclass, field + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, nn +from torch.nn.attention import SDPBackend, sdpa_kernel +from torch.utils.checkpoint import checkpoint + +from toolkit.models.v2._mixin import OstrisModelMixin + +try: + from flash_attn_interface import flash_attn_func + + _FLASH_ATTN_AVAILABLE = True +except ImportError: + flash_attn_func = None + _FLASH_ATTN_AVAILABLE = False + +ATTENTION_BACKENDS = ("native", "cudnn", "flash") + + +@dataclass +class Flux3TransformerParams: + # content modes -> latent channels (text is separate: context_in_dim) + in_channels: dict = field( + default_factory=lambda: {"image": 128, "image_cond": 128} + ) + # forward kwarg name -> content mode + sequence: dict = field( + default_factory=lambda: {"x_image": "image", "x_image_cond": "image_cond"} + ) + # pooled-vector embedder width; None = no vector_in + vec_in_dim: int | None = None + context_in_dim: int = 20480 + hidden_size: int = 3072 + num_heads: int = 24 + depth: int = 5 + depth_single_blocks: int = 28 + axes_dim: list[int] = field(default_factory=lambda: [32, 32, 32, 32]) + theta: int = 10000 + mlp_ratio: float = 3.0 + qkv_bias: bool = False + gate_type: str | None = None + attn_mode: str = "native" + # modes with a final_layer (velocity outputs); None = every content mode + out_modes: list[str] | None = field(default_factory=lambda: ["image"]) + + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + assert dim % 2 == 0 + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack( + [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1 + ) + out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2) + return out.float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) + + +class Attention(nn.Module): + def forward( + self, q: Tensor, k: Tensor, v: Tensor, attn_mode: str = "native" + ) -> Tensor: + if attn_mode == "flash": + if flash_attn_func is None: + raise RuntimeError("flash_attn_func is unavailable") + b, h, n, d = q.shape + q, k, v = (rearrange(t, "b h n d -> b n h d") for t in (q, k, v)) + x = flash_attn_func(q, k, v)[0] + return x.reshape(b, n, h * d) + if attn_mode == "cudnn": + with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + x = F.scaled_dot_product_attention(q, k, v) + else: # native + x = F.scaled_dot_product_attention(q, k, v) + return rearrange(x, "b h n d -> b n (h d)") + + +def _bf16_if_required(module: nn.Module, x: Tensor) -> Tensor: + if getattr(module, "_requires_bf16_input", False) and x.dtype != torch.bfloat16: + return x.to(torch.bfloat16) + return x + + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + emb = torch.cat( + [ + rope(ids[..., i], self.axes_dim[i], self.theta) + for i in range(len(self.axes_dim)) + ], + dim=-3, + ) + return emb.unsqueeze(1) + + +def timestep_embedding( + t: Tensor, + dim: int, + max_period: int = 10000, + time_factor: float = 1000.0, + out_dtype: torch.dtype | None = None, +) -> Tensor: + t = time_factor * t + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(0, half, device=t.device, dtype=torch.float32) + / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(out_dtype if out_dtype is not None else t.dtype) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int, disable_bias: bool = False): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=not disable_bias) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=not disable_bias) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=dtype) * self.scale + + +class QKNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + return self.query_norm(q).to(v), self.key_norm(k).to(v) + + +class SiLUActivation(nn.Module): + def __init__(self): + super().__init__() + self.gate_fn = nn.SiLU() + + def forward(self, x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return self.gate_fn(x1) * x2 + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool, disable_bias: bool = False): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=not disable_bias) + + def forward(self, vec: Tensor): + out = self.lin(F.silu(vec)) + if out.ndim == 2: + out = out[:, None, :] + out = out.chunk(self.multiplier, dim=-1) + return out[:3], out[3:] if self.is_double else None + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, out_channels, bias=False) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=False) + ) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + if vec.ndim == 2: + vec = vec[:, None, :] + s = self.adaLN_modulation[0](vec) # type: ignore + shift_w, scale_w = self.adaLN_modulation[1].weight.chunk(2) + + # x = (1 + scale) * self.norm_final(x) + shift + # avoid materializing the full (seq, 2*hidden) projection to reduce memory + x = self.norm_final(x) + x = x * (F.linear(s, scale_w) + 1) + x = x + F.linear(s, shift_w) + return self.linear(x) + + +class ModeBlock(nn.Module): + """Per-mode (early) block: self attention within one stream.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 3.0, + qkv_bias: bool = False, + gate_type: str | None = None, + use_gated_mlp: bool = True, + ): + super().__init__() + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.hidden_size = hidden_size + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.mlp_mult_factor = 2 if use_gated_mlp else 1 + self.gate_type = gate_type + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.mlp_in = nn.Linear( + hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor, bias=False + ) + if gate_type is not None: + self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False) + + self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False) + self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) + + self.norm = QKNorm(head_dim) + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp_act = ( + SiLUActivation() if use_gated_mlp else nn.GELU(approximate="tanh") + ) + self.attention = Attention() + self.attn_mode = "native" + + def forward(self, x: Tensor, pe: Tensor, mod: tuple[Tensor, ...]) -> Tensor: + mod_shift, mod_scale, mod_gate = mod + x_mod = (1 + mod_scale) * self.pre_norm(x) + mod_shift + + q = rearrange( + self.q_proj(_bf16_if_required(self.q_proj, x_mod)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + k = rearrange( + self.k_proj(_bf16_if_required(self.k_proj, x_mod)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + v = rearrange( + self.v_proj(_bf16_if_required(self.v_proj, x_mod)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + mlp = self.mlp_in(_bf16_if_required(self.mlp_in, x_mod)) + + q, k = self.norm(q, k, v) + q, k = apply_rope(q, k, pe) + attn = self.attention(q, k, v, self.attn_mode) + + output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out( + _bf16_if_required(self.mlp_out, self.mlp_act(mlp)) + ) + if self.gate_type is not None: + output = output * F.silu(self.gate_proj(x_mod)) + return x + mod_gate * output + + +class SingleStreamBlock(nn.Module): + """Joint block: text + all active content modes in one attention sequence.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 3.0, + gate_type: str | None = None, + use_gated_mlp: bool = True, + ): + super().__init__() + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.hidden_size = hidden_size + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.mlp_mult_factor = 2 if use_gated_mlp else 1 + self.gate_type = gate_type + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.mlp_in = nn.Linear( + hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor, bias=False + ) + self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False) + self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) + if gate_type is not None: + self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False) + + self.norm = QKNorm(head_dim) + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp_act = ( + SiLUActivation() if use_gated_mlp else nn.GELU(approximate="tanh") + ) + self.attention = Attention() + self.attn_mode = "native" + + def forward( + self, + x: Tensor, + ctx: Tensor, + pe: Tensor, + mod_x: tuple[Tensor, ...], + mod_ctx: tuple[Tensor, ...], + ) -> tuple[Tensor, Tensor]: + mod_shift_x, mod_scale_x, mod_gate_x = mod_x + mod_shift_ctx, mod_scale_ctx, mod_gate_ctx = mod_ctx + + x_mod = (1 + mod_scale_x) * self.pre_norm(x) + mod_shift_x + ctx_mod = (1 + mod_scale_ctx) * self.pre_norm(ctx) + mod_shift_ctx + combined = torch.cat((ctx_mod, x_mod), dim=1) + + q = rearrange( + self.q_proj(_bf16_if_required(self.q_proj, combined)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + k = rearrange( + self.k_proj(_bf16_if_required(self.k_proj, combined)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + v = rearrange( + self.v_proj(_bf16_if_required(self.v_proj, combined)), + "B L (H D) -> B H L D", + H=self.num_heads, + ) + mlp = self.mlp_in(_bf16_if_required(self.mlp_in, combined)) + + q, k = self.norm(q, k, v) + q, k = apply_rope(q, k, pe) + attn = self.attention(q, k, v, self.attn_mode) + + output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out( + _bf16_if_required(self.mlp_out, self.mlp_act(mlp)) + ) + if self.gate_type is not None: + output = output * F.silu(self.gate_proj(combined)) + + ctx_out, x_out = torch.split(output, [ctx.shape[1], x.shape[1]], dim=1) + return ctx + mod_gate_ctx * ctx_out, x + mod_gate_x * x_out + + +class Flux3Transformer(nn.Module, OstrisModelMixin): + @classmethod + def get_transformer_block_names(cls): + return [ + "content_mode_blocks.image", + "content_mode_blocks.image_cond", + "txt_mode_blocks", + "single_blocks", + ] + + def get_offload_ignore_modules(self): + # modulation tables stay resident on the compute device; LastLayer + # reads adaLN weights raw (bypasses offload forward hijacks) + return [m for m in self.modules() if isinstance(m, (Modulation, LastLayer))] + + def __init__(self, params: Flux3TransformerParams | None = None): + super().__init__() + if params is None: + params = Flux3TransformerParams() + self.params = params + + self.sequence = params.sequence + for _, v in params.sequence.items(): + assert v in params.in_channels, (params.sequence, params.in_channels) + + self.in_channels = params.in_channels + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.depth = params.depth + self.depth_single_blocks = params.depth_single_blocks + self.mlp_ratio = params.mlp_ratio + self.qkv_bias = params.qkv_bias + self.gate_type = params.gate_type + self.attn_mode = params.attn_mode + + self.gradient_checkpointing = False + + if params.hidden_size % params.num_heads != 0: + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by " + f"num_heads {params.num_heads}" + ) + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError( + f"Got {params.axes_dim} but expected positional dim {pe_dim}" + ) + + self.pe_embedder = EmbedND( + dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim + ) + + self.emb_in = nn.ModuleDict( + { + m: nn.Linear(c, params.hidden_size, bias=False) + for m, c in params.in_channels.items() + } + ) + self.txt_in = nn.Linear(params.context_in_dim, params.hidden_size, bias=False) + + self.time_in = MLPEmbedder( + in_dim=256, hidden_dim=params.hidden_size, disable_bias=True + ) + self.vector_in: nn.Module = ( + MLPEmbedder(params.vec_in_dim, params.hidden_size, disable_bias=True) + if params.vec_in_dim is not None + else nn.Identity() + ) + + content_modes = sorted(params.in_channels.keys()) + self.early_stream_modulations = nn.ModuleDict( + { + **{ + m: Modulation(params.hidden_size, double=False, disable_bias=True) + for m in content_modes + }, + "txt": Modulation(params.hidden_size, double=False, disable_bias=True), + } + ) + self.single_stream_modulations = nn.ModuleDict( + { + **{ + m: Modulation(params.hidden_size, double=False, disable_bias=True) + for m in content_modes + }, + "txt": Modulation(params.hidden_size, double=False, disable_bias=True), + } + ) + + self.content_mode_blocks = nn.ModuleDict( + { + m: nn.ModuleList( + [ + ModeBlock( + params.hidden_size, + params.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + gate_type=params.gate_type, + ) + for _ in range(params.depth) + ] + ) + for m in content_modes + } + ) + self.txt_mode_blocks = nn.ModuleList( + [ + ModeBlock( + params.hidden_size, + params.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + gate_type=params.gate_type, + ) + for _ in range(params.depth) + ] + ) + + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock( + params.hidden_size, + params.num_heads, + mlp_ratio=params.mlp_ratio, + gate_type=params.gate_type, + ) + for _ in range(params.depth_single_blocks) + ] + ) + + out_modes = ( + list(params.in_channels) if params.out_modes is None else params.out_modes + ) + assert set(out_modes) <= set(params.in_channels), (out_modes, params.in_channels) + self.final_layer = nn.ModuleDict( + { + m: LastLayer(params.hidden_size, c) + for m, c in params.in_channels.items() + if m in out_modes + } + ) + + self.set_attention_backend(params.attn_mode) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def enable_gradient_checkpointing(self, enable: bool = True): + self.gradient_checkpointing = enable + + def disable_gradient_checkpointing(self): + self.gradient_checkpointing = False + + def set_attention_backend(self, backend: str): + if backend not in ATTENTION_BACKENDS: + raise ValueError( + f"Unknown attention backend {backend!r}; choose from {ATTENTION_BACKENDS}" + ) + if backend == "flash" and not _FLASH_ATTN_AVAILABLE: + raise RuntimeError( + "attention backend 'flash' requested but flash_attn_interface " + "(flash-attn 3) is not installed" + ) + if backend == "cudnn" and not torch.cuda.is_available(): + raise RuntimeError("attention backend 'cudnn' requires CUDA") + self.attn_mode = backend + for module in self.modules(): + if hasattr(module, "attn_mode"): + module.attn_mode = backend + + def _ckpt(self, module, *args): + if self.gradient_checkpointing and torch.is_grad_enabled(): + return checkpoint(module, *args, use_reentrant=False) + return module(*args) + + def forward( + self, + ctx: Tensor, + ctx_ids: Tensor, + timesteps_ctx: Tensor, + vector: Tensor | None = None, + **kwargs: Tensor, + ) -> dict[str, Tensor]: + batch_size = ctx.shape[0] + + x_by_mod: dict[str, list] = defaultdict(list) + ids_by_mod: dict[str, list] = defaultdict(list) + ts_by_mod: dict[str, list] = defaultdict(list) + lens_by_mod: dict[str, list] = defaultdict(list) + names_by_mod: dict[str, list] = defaultdict(list) + for k, v in self.sequence.items(): + if k not in kwargs or kwargs[k] is None: + continue + assert k + "_ids" in kwargs and k + "_timesteps" in kwargs + x_by_mod[v].append(kwargs[k]) + ids_by_mod[v].append(kwargs[k + "_ids"]) + ts_by_mod[v].append(kwargs[k + "_timesteps"]) + lens_by_mod[v].append(kwargs[k].shape[1]) + names_by_mod[v].append(k) + + active = list(x_by_mod.keys()) + all_mods = sorted(self.in_channels.keys()) + inactive = [m for m in all_mods if m not in active] + dummy_len = 2 # attention kernels need >=2 + + device = x_by_mod[active[0]][0].device + dtype = x_by_mod[active[0]][0].dtype + for k in all_mods: + if k in active: + x_by_mod[k] = torch.cat(x_by_mod[k], 1) + ids_by_mod[k] = torch.cat(ids_by_mod[k], 1) + ts_by_mod[k] = torch.cat(ts_by_mod[k], 1) + else: + x_by_mod[k] = torch.zeros( + batch_size, + dummy_len, + self.in_channels[k], + device=device, + dtype=dtype, + ) + ids_by_mod[k] = torch.zeros( + batch_size, dummy_len, 4, device=device, dtype=torch.int64 + ) + ts_by_mod[k] = torch.zeros( + batch_size, dummy_len, device=device, dtype=dtype + ) + + vec_by_mod: dict[str, Tensor] = {} + for k in all_mods: + ts = rearrange(ts_by_mod[k], "b t -> (b t)") + vec_by_mod[k] = rearrange( + self.time_in(timestep_embedding(ts, 256, out_dtype=dtype)), + "(b t) d -> b t d", + b=batch_size, + ) + ts_ctx = rearrange(timesteps_ctx, "b t -> (b t)") + vec_ctx = rearrange( + self.time_in(timestep_embedding(ts_ctx, 256, out_dtype=dtype)), + "(b t) d -> b t d", + b=batch_size, + ) + + if self.params.vec_in_dim is not None: + assert vector is not None and vector.ndim == 2 + vector_emb = self.vector_in(vector) + for k in all_mods: + vec_by_mod[k] = vec_by_mod[k] + vector_emb[:, None, :] + vec_ctx = vec_ctx + vector_emb[:, None, :] + + early_content_mods: dict[str, tuple] = {} + for key in self.early_stream_modulations: + if key == "txt": + continue + mod, _ = self.early_stream_modulations[key](vec_by_mod[key]) + early_content_mods[key] = mod + early_txt_mod, _ = self.early_stream_modulations["txt"](vec_ctx) + + single_content_mods: dict[str, list] = {} + for key in self.single_stream_modulations: + if key == "txt": + continue + mod, _ = self.single_stream_modulations[key](vec_by_mod[key]) + single_content_mods[key] = list(mod) + single_txt_mod, _ = self.single_stream_modulations["txt"](vec_ctx) + + content_by_mod = {k: self.emb_in[k](x_by_mod[k]) for k in all_mods} + txt = self.txt_in(ctx) + + pe_by_mod = {k: self.pe_embedder(ids_by_mod[k]) for k in all_mods} + pe_ctx = self.pe_embedder(ctx_ids) + + for d in range(self.depth): + for m in all_mods: + content_by_mod[m] = self._ckpt( + self.content_mode_blocks[m][d], + content_by_mod[m], + pe_by_mod[m], + early_content_mods[m], + ) + txt = self._ckpt(self.txt_mode_blocks[d], txt, pe_ctx, early_txt_mod) + + img = torch.cat([content_by_mod[k] for k in active], dim=1) + content_mod = tuple( + torch.cat([single_content_mods[k][i] for k in active], dim=1) + for i in range(3) + ) + pe_content = torch.cat([pe_by_mod[k] for k in active], dim=2) + pe = torch.cat((pe_ctx, pe_content), dim=2) + + for block in self.single_blocks: + txt, img = self._ckpt( + block, + img, + txt, + pe, + content_mod, + single_txt_mod, + ) + + img_by_mod: dict[str, Tensor] = {} + pos = 0 + for k in active: + length = sum(lens_by_mod[k]) + img_by_mod[k] = img[:, pos : pos + length] + pos += length + for k in inactive: + img_by_mod[k] = content_by_mod[k] + + final = { + k: self._ckpt(self.final_layer[k], img_by_mod[k], vec_by_mod[k]) + for k in self.final_layer + } + out: dict[str, Tensor] = {} + for k in active: + if k not in final: + continue + pos = 0 + for name, length in zip(names_by_mod[k], lens_by_mod[k]): + out[name] = final[k][:, pos : pos + length] + pos += length + return out diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 9cee405..fe8c1fe 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1662,6 +1662,38 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "datasets.auto_frame_count", ], }, + { + name: "flux3", + label: "FLUX.3 (small)", + group: "image", + gateUrl: "https://huggingface.co/herzogenhorn/f3-image-small", + defaults: { + // default updates when [selected, unselected] in the UI + "config.process[0].model.name_or_path": [ + "herzogenhorn/f3-image-small", + defaultNameOrPath, + ], + "config.process[0].model.quantize": [true, false], + "config.process[0].model.qtype": ["convrot8", "qfloat8"], + "config.process[0].model.quantize_te": [true, false], + "config.process[0].model.qtype_te": ["convrot8", "qfloat8"], + "config.process[0].model.low_vram": [true, false], + "config.process[0].sample.sampler": ["flowmatch", "flowmatch"], + "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], + "config.process[0].sample.width": [1024, 1024], + "config.process[0].sample.height": [1024, 1024], + "config.process[0].sample.guidance_scale": [3.5, 4], + "config.process[0].sample.sample_steps": [28, 25], + "config.process[0].train.timestep_type": ["shift", "sigmoid"], + }, + disableSections: ["network.conv"], + additionalSections: [ + "datasets.control_path", + "sample.ctrl_img", + "model.layer_offloading", + "model.low_vram", + ], + }, { name: "flux2_klein_4b", label: "FLUX.2-klein-base-4B", diff --git a/extensions_built_in/inference_engine/latent_format_tables.py b/extensions_built_in/inference_engine/latent_format_tables.py index b02bce4..106ce71 100644 --- a/extensions_built_in/inference_engine/latent_format_tables.py +++ b/extensions_built_in/inference_engine/latent_format_tables.py @@ -540,4 +540,45 @@ LATENT_FORMATS = { [0.0447, 0.0994, 0.0881], ], }, + "Flux3Image": { + # FLUX.3 image VAE (FLUX.2 packed f8 autoencoder, BN-normalized): 32 ch x 2x2 shuffle. + # least-squares fit on flux3-small latents (2026-09-12), r2 ~ [0.74, 0.76, 0.76] + "channels": 128, "dims": 2, "spatial": 16, "temporal": 1, + "reshape": "flux2_2x2", + "bias": [-0.01281, -0.10207, -0.2469], + "factors": [ + [0.01443, 0.01462, 0.02385], + [-0.01563, -0.05629, 0.03518], + [-0.00366, 0.01501, 0.0998], + [0.17186, 0.27951, 0.40425], + [0.05145, -0.01879, -0.02143], + [0.01929, -0.00069, -0.03722], + [-0.0942, 0.11153, -0.09854], + [0.03436, 0.01086, -0.00353], + [-0.33485, -0.22427, -0.04648], + [0.01356, -0.00215, 0.03131], + [-0.00078, 0.0166, -0.01035], + [0.05476, 0.04449, -0.06473], + [0.00259, -0.00427, -0.01085], + [0.02877, 0.00309, -0.02534], + [-0.01333, 0.0106, 0.01833], + [-0.0832, 0.09934, 0.07942], + [-0.00951, -0.00139, 0.0027], + [0.01677, -0.01178, -0.02627], + [-0.01436, -0.02519, -0.01887], + [0.13145, -0.0941, 0.02256], + [-0.0128, -0.00109, 0.01468], + [0.03173, -0.01365, 0.01746], + [-0.05829, 0.01014, -0.02845], + [-0.02415, 0.01729, 0.03453], + [0.0132, 0.03026, 0.00655], + [0.004, 0.00867, -0.01574], + [0.07235, -0.00765, 0.03589], + [0.05798, 0.01141, 0.00962], + [-0.01311, -0.02256, -0.04102], + [-0.03182, -0.02579, 0.00077], + [0.02911, 0.01766, -0.01212], + [0.01937, -0.00168, -0.02601], + ], + }, } diff --git a/extensions_built_in/inference_engine/latent_preview.py b/extensions_built_in/inference_engine/latent_preview.py index a29f1ed..5979f4d 100644 --- a/extensions_built_in/inference_engine/latent_preview.py +++ b/extensions_built_in/inference_engine/latent_preview.py @@ -50,6 +50,7 @@ ARCH_LATENT_FORMAT = { "anima": "Wan21", "cogview4": "Flux", "prx_pixel": None, # pixel space + "flux3": "Flux3Image", } diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index a2357c2..c77807b 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -54,6 +54,7 @@ ARCH_REGISTRY: Dict[str, dict] = { "wan22_14b_i2v": {"modality": "video", "model": {"name_or_path": "ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "quantize": True, "quantize_te": True, "low_vram": True}, "sample": {**VID, "guidance_scale": 3.5}, "needs_control_image": True}, "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, + "flux3": {"modality": "image", "model": {"name_or_path": "herzogenhorn/f3-image-small", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}}, "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}}, "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}}, "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},