From bcb89cd3bfdbaa102d7c489094754c1b085e094b Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Fri, 16 Oct 2026 08:23:49 +0800 Subject: [PATCH 1/8] [model,test,docs] feat: add FLUX.3 adapter --- .agents/knowledge/architecture.md | 1 + README.md | 16 +- diffusers | 2 +- examples/README.md | 8 + examples/grpo/lora/flux3/default.yaml | 113 ++ guidance/datasets.md | 2 +- src/flow_factory/hparams/model_args.py | 1 + src/flow_factory/models/flux/_output.py | 192 +- src/flow_factory/models/flux/flux3.py | 1563 +++++++++++++++++ src/flow_factory/models/registry.py | 1 + tests/docs/test_flux3_docs.py | 29 + tests/models/test_all_model_loaders.py | 34 + tests/models/test_flux3.py | 712 ++++++++ .../models/test_modern_image_output_codecs.py | 9 + .../models/test_offline_training_guidance.py | 2 + 15 files changed, 2678 insertions(+), 7 deletions(-) create mode 100644 examples/grpo/lora/flux3/default.yaml create mode 100644 src/flow_factory/models/flux/flux3.py create mode 100644 tests/docs/test_flux3_docs.py create mode 100644 tests/models/test_flux3.py diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md index 6a855052..0d2bfb77 100644 --- a/.agents/knowledge/architecture.md +++ b/.agents/knowledge/architecture.md @@ -130,6 +130,7 @@ extensions are `GRPOGuardTrainer → GRPOTrainer`, `DPPOTrainer → GRPOTrainer` | `flux1-kontext` | `Flux1KontextAdapter` | Image-to-Image | | `flux2` | `Flux2Adapter` | Text-to-Image & Image(s)-to-Image | | `flux2-klein` | `Flux2KleinAdapter` | Text-to-Image & Image(s)-to-Image | +| `flux3` | `Flux3Adapter` | Text-to-Image & Image(s)-to-Image | | `qwen-image` | `QwenImageAdapter` | Text-to-Image | | `qwen-image-edit-plus` | `QwenImageEditPlusAdapter` | Image(s)-to-Image | | `qwen-image-2.1` | `QwenImage21Adapter` | Text-to-Image & Image(s)-to-Image | diff --git a/README.md b/README.md index 02fc59eb..17501764 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,16 @@ # 🔥 News +* **[2026-10-16]** **FLUX.3 Image Klein** support! Train both the base (`black-forest-labs/flux-3-image-klein-base`) + and step-distilled flow-map (`black-forest-labs/flux-3-image-klein`) text-to-image and multi-reference + image-to-image models with the [GRPO + LoRA recipe](examples/grpo/lora/flux3/default.yaml). + Until FLUX.3 reaches a released Diffusers package, install the pinned submodule: +```bash +git submodule update --init +pip install -e ./diffusers +pip install -e . +``` + * **[2026-09-20]** **Qwen-Image 2.1** support! Train the unified text-to-image and ordered multi-image editing model with the [GRPO + LoRA recipe](examples/grpo/lora/qwen_image_2_1/default.yaml). Until Qwen-Image 2.1 reaches a released Diffusers package, install the pinned submodule: @@ -74,11 +84,13 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac Image(s)-to-ImageQwen-Image-Edit-250920Bqwen-image-edit-plus Qwen-Image-Edit-251120Bqwen-image-edit-plus - Text-to-Image & Image(s)-to-ImageFLUX.2-dev32Bflux2 + Text-to-Image & Image(s)-to-ImageFLUX.2-dev32Bflux2 FLUX.2-klein-4B4Bflux2-klein FLUX.2-klein-9B9Bflux2-klein FLUX.2-klein-base-4B4Bflux2-klein FLUX.2-klein-base-9B9Bflux2-klein + flux-3-image-klein-base5.5Bflux3 + flux-3-image-klein5.5Bflux3 BAGEL-7B-MoT14Bbagel SenseNova-U1 1.016Bsensenova SenseNova-U1 1.516Bsensenova @@ -107,7 +119,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac > To support new models, see [Guidance/New Model](guidance/new_model.md). > **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`, -> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, +> `flux1-kontext`, `flux2`, `flux2-klein`, `flux3`, `qwen-image`, `qwen-image-edit-plus`, > `qwen-image-2.1`, `z-image`, > `bagel`, `sensenova`, `wan2_t2v`, `wan2_i2v`, `ltx2_t2av`, `ltx2_i2av`, and all > MiniMax H3 workflows. Video/audio targets are encoded on demand and are never written to diff --git a/diffusers b/diffusers index 80c7ed26..7ecb1e1b 160000 --- a/diffusers +++ b/diffusers @@ -1 +1 @@ -Subproject commit 80c7ed262aeffbeb43ef13ae04baeb9b84515a69 +Subproject commit 7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941 diff --git a/examples/README.md b/examples/README.md index f0d98462..54815474 100644 --- a/examples/README.md +++ b/examples/README.md @@ -96,6 +96,14 @@ The same adapter also accepts prompt-only records. It currently requires the bun submodule, enforces `per_device_batch_size: 1`, and keeps prefix KV caching disabled for exact rollout/replay parity. +## FLUX.3 example + +- [`flux3` text-to-image & multi-reference editing GRPO](grpo/lora/flux3/default.yaml) + +Supports both `black-forest-labs/flux-3-image-klein-base` (true CFG) and +`black-forest-labs/flux-3-image-klein` (step-distilled flow map, no CFG) via the bundled +Diffusers submodule. + ## MiniMax H3 validation status The T2VA `debug.yaml` recipe is real-weight validated with the 61 GB checkpoint diff --git a/examples/grpo/lora/flux3/default.yaml b/examples/grpo/lora/flux3/default.yaml new file mode 100644 index 00000000..adac52ff --- /dev/null +++ b/examples/grpo/lora/flux3/default.yaml @@ -0,0 +1,113 @@ +# FLUX.3 currently requires the bundled Diffusers submodule: +# git submodule update --init +# pip install -e ./diffusers + +# Environment Configuration +launcher: "accelerate" +config_file: config/deepspeed/deepspeed_zero2.yaml # Switch to config/accelerate_configs/fsdp2.yaml if you have limited GPU memory. +num_processes: 8 +main_process_port: 29500 +mixed_precision: "bf16" + +# Data Configuration +data: + datasets: + - name: default + dataset_dir: "dataset/sharegpt4o_image_mini" + train: { weight: 1 } + eval: {} + preprocessing_batch_size: 8 + dataloader_num_workers: 16 + force_reprocess: true + cache_dir: "~/.cache/flow_factory/datasets" + sampler_type: "auto" + sampler_subgroup_size: null + +# Model Configuration +model: + finetune_type: "lora" + lora_rank: 64 + lora_alpha: 128 + target_modules: "default" + model_name_or_path: "black-forest-labs/flux-3-image-klein-base" # Options: black-forest-labs/flux-3-image-klein-base (CFG), black-forest-labs/flux-3-image-klein (step-distilled flow-map) + model_type: "flux3" + resume_path: null + resume_type: null + +log: + run_name: null + project: "Flow-Factory" + logging_backend: "wandb" + save_dir: "saves/" + save_freq: 20 + save_model_only: true + +# Training Configuration +train: + trainer_type: "grpo" + reward_optimization_overlap: false + reward_optimization_overlap_mode: "ready" + reward_optimization_overlap_poll_interval: 0.05 + advantage_aggregation: "gdpo" + clip_range: 1.0e-4 + adv_clip_range: 5.0 + kl_type: "v-based" + kl_beta: 0 + ref_param_device: "cuda" + + # Sampling settings + resolution: 384 + condition_image_size: 384 + num_inference_steps: 10 # Use 4-16 with guidance_scale: 1.0 for black-forest-labs/flux-3-image-klein + guidance_scale: 3.5 + + # Batch and sampling + per_device_batch_size: 1 + group_size: 16 + global_std: false + unique_sample_num_per_epoch: 48 + gradient_step_per_epoch: 2 + gradient_accumulation_steps: auto + + ema_decay: 0.9 + ema_update_interval: 4 + ema_device: "cuda" + enable_gradient_checkpointing: true + offload_samples_to_cpu: false + seed: 42 + +# Scheduler Configuration +scheduler: + dynamics_type: "Flow-SDE" + noise_level: 0.7 + num_sde_steps: 1 + sde_steps: [1, 2, 3] + seed: 42 + +# Evaluation Configuration +eval: + resolution: 1024 + condition_image_size: 1024 + per_device_batch_size: 1 + guidance_scale: 3.5 + num_inference_steps: 28 + eval_freq: 20 + seed: 42 + +# Reward Model Configuration +rewards: + - name: "pick_score" + reward_model: "PickScore" + weight: 1 + batch_size: 16 + device: "cuda" + dtype: bfloat16 + +# Optimizer Configuration +optimizers: + - name: default + learning_rate: 3.0e-4 + weight_decay: 1.0e-4 + betas: [0.9, 0.999] + eps: 1.0e-8 + max_grad_norm: 1.0 diff --git a/guidance/datasets.md b/guidance/datasets.md index ab3aaf35..726fb708 100644 --- a/guidance/datasets.md +++ b/guidance/datasets.md @@ -365,7 +365,7 @@ output semantics. | Offline status | Model types | Notes | |---|---|---| -| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `qwen-image-2.1`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. Qwen-Image 2.1 keeps ordered condition blocks before the target block and runs offline batches one sample at a time. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. | +| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `flux3`, `qwen-image`, `qwen-image-edit-plus`, `qwen-image-2.1`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. Qwen-Image 2.1 keeps ordered condition blocks before the target block and runs offline batches one sample at a time. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. | | Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. | | Supported | `wan2_i2v` | Input media uses a checkpoint-specific contract. Expanded-timestep TI2V and standard Wan2.1 I2V accept exactly `first_frame`; dedicated Wan2.1 FLF2V requires both `first_frame` and `last_frame`; Wan2.2 I2V-A14B accepts `first_frame` plus an optional VAE-only `last_frame`. Condition pixels are cached at configured geometry and encoded with VAE posterior mode once per batch. Video targets require `fps`; offline execution is B=1. | | Supported | `ltx2_t2av`, `ltx2_i2av` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. Both streams are aligned to the official LTX2 clock and encoded/packed on the fly. I2AV requires the `first_frame` image slot, substitutes its posterior-mode first latent into each target, and excludes the pinned tokens with an active mask. | diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py index e022ed11..fdb4412a 100644 --- a/src/flow_factory/hparams/model_args.py +++ b/src/flow_factory/hparams/model_args.py @@ -165,6 +165,7 @@ class ModelArguments(ArgABC): "flux1-kontext", "flux2", "flux2-klein", + "flux3", "qwen-image", "qwen-image-edit-plus", "qwen-image-2.1", diff --git a/src/flow_factory/models/flux/_output.py b/src/flow_factory/models/flux/_output.py index 3ecf69e9..9eb4db6e 100644 --- a/src/flow_factory/models/flux/_output.py +++ b/src/flow_factory/models/flux/_output.py @@ -105,6 +105,91 @@ def prepare_flux1_output_latents( return packed, target_ids +def patchify_flux2_latents(latents: torch.Tensor) -> torch.Tensor: + """Convert 4D VAE latents ``(B, C, H, W)`` to 2x2 patchified ``(B, C*4, H//2, W//2)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for 2x2 patchify, received shape {tuple(latents.shape)}" + ) + batch_size, channels, height, width = latents.shape + if height % 2 or width % 2: + raise ValueError( + f"expected even latent height/width for 2x2 patchify, received {(height, width)}" + ) + latents = latents.view(batch_size, channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + return latents.reshape(batch_size, channels * 4, height // 2, width // 2) + + +def unpatchify_flux2_latents(latents: torch.Tensor) -> torch.Tensor: + """Convert 2x2 patchified latents ``(B, C*4, H, W)`` back to ``(B, C, H*2, W*2)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for 2x2 unpatchify, received shape {tuple(latents.shape)}" + ) + batch_size, channels, height, width = latents.shape + if channels % 4: + raise ValueError( + f"expected channel count divisible by 4 for 2x2 unpatchify, received {channels}" + ) + latents = latents.reshape(batch_size, channels // 4, 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + return latents.reshape(batch_size, channels // 4, height * 2, width * 2) + + +def pack_flux3_latents(latents: torch.Tensor) -> torch.Tensor: + """Pack 4D patchified latents ``(B, C, H, W)`` into tokens ``(B, H*W, C)``.""" + if latents.ndim != 4: + raise ValueError( + f"expected 4D BCHW latents for FLUX.3 packing, received shape {tuple(latents.shape)}" + ) + return latents.flatten(2).transpose(1, 2) + + +def unpack_flux3_latents( + latents: torch.Tensor, + latent_height: int, + latent_width: int, +) -> torch.Tensor: + """Unpack FLUX.3 latent tokens ``(B, H*W, C)`` onto a 2D grid ``(B, C, H, W)``.""" + if latents.ndim != 3: + raise ValueError( + f"expected 3D (B, seq_len, C) latents for FLUX.3 unpacking, " + f"received shape {tuple(latents.shape)}" + ) + batch_size, seq_len, channels = latents.shape + if latent_height <= 0 or latent_width <= 0 or latent_height * latent_width != seq_len: + raise ValueError( + f"FLUX.3 latent token count {seq_len} does not match " + f"latent grid {latent_height} x {latent_width}" + ) + return latents.view(batch_size, latent_height, latent_width, channels).permute(0, 3, 1, 2) + + +def prepare_flux3_image_ids( + latent_height: int, + latent_width: int, + *, + image_rope_scale: float = 0.5, + l_coord: int = 0, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Build 4D position IDs ``(t, h, w, l)`` of shape ``(H*W, 4)`` for FLUX.3 image tokens.""" + if latent_height <= 0 or latent_width <= 0: + raise ValueError( + f"expected positive FLUX.3 latent grid dimensions, received {(latent_height, latent_width)}" + ) + t_coords = torch.arange(1) + h_coords = torch.arange(latent_height) + w_coords = torch.arange(latent_width) + l_coords = torch.tensor([l_coord]) + ids = torch.cartesian_prod(t_coords, h_coords, w_coords, l_coords).to(dtype=torch.float32) + ids[:, 1:3] = ids[:, 1:3] * float(image_rope_scale) + if device is not None: + ids = ids.to(device) + return ids + + def encode_flux2_output_images( adapter: Any, pixel_values: torch.Tensor, @@ -135,6 +220,41 @@ def encode_flux2_output_images( ) +def encode_flux3_output_images( + adapter: Any, + pixel_values: torch.Tensor, + generator: Optional[torch.Generator], + *, + image_rope_scale: float = 0.5, +) -> EncodedImageTensor: + """Sample the FLUX.3 target posterior, normalize, and pack tokens with scaled RoPE IDs.""" + latents = encode_flux2_vae_image( + adapter, + pixel_values, + sample_mode="sample", + generator=generator, + ) + batch_size = latents.shape[0] + latent_height, latent_width = latents.shape[-2:] + latent_ids = ( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=image_rope_scale, + l_coord=0, + device=adapter.device, + ) + .unsqueeze(0) + .repeat(batch_size, 1, 1) + ) + packed = pack_flux3_latents(latents).to(adapter.device) + return EncodedImageTensor( + latents=packed, + forward_context={"latent_ids": latent_ids}, + decode_context={"latent_ids": latent_ids}, + ) + + def encode_flux2_vae_image( adapter: Any, pixel_values: torch.Tensor, @@ -142,10 +262,10 @@ def encode_flux2_vae_image( sample_mode: VAELatentSampleMode, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: - """Apply the exact FLUX.2 transform shared by condition and target roles. + """Apply the exact FLUX.2/FLUX.3 transform shared by condition and target roles. Args: - adapter: FLUX.2 adapter exposing the canonical VAE and patchify primitive. + adapter: FLUX.2 or FLUX.3 adapter exposing the canonical VAE and patchify primitive. pixel_values: Preprocessed BCHW image tensor. sample_mode: Explicit posterior selection for the calling pipeline role. generator: Generator forwarded unchanged for posterior sampling. @@ -165,7 +285,11 @@ def encode_flux2_vae_image( generator=generator, source=f"{type(adapter).__name__} FLUX.2 VAE encode", ) - latents = adapter.pipeline._patchify_latents(latents) + pipeline_patchify = getattr(getattr(adapter, "pipeline", None), "_patchify_latents", None) + if callable(pipeline_patchify): + latents = pipeline_patchify(latents) + else: + latents = patchify_flux2_latents(latents) if latents.ndim != 4: raise ValueError( f"{type(adapter).__name__} FLUX.2 patchify expected BCHW output, " @@ -259,10 +383,72 @@ def prepare_flux2_condition_latents( return packed, image_latent_ids +def prepare_flux3_condition_latents( + adapter: Any, + images: List[torch.Tensor], + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + image_rope_scale: float = 0.5, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compose FLUX.3 reference-image tokens and ``l = -(1 + idx)`` position IDs. + + Args: + adapter: FLUX.3 adapter exposing the canonical VAE. + images: One or more preprocessed BCHW condition tensors for one sample. + batch_size: Number of prompt rows that reuse the condition sequence. + device: Target device for condition tensors and identifiers. + dtype: Target compute dtype for packed condition tokens. + image_rope_scale: Spatial RoPE coordinate multiplier (0.5 in FLUX.3). + + Returns: + Packed condition latents ``(B, cond_seq_len, C)`` and position IDs ``(B, cond_seq_len, 4)``. + """ + if not images: + raise ValueError(f"{type(adapter).__name__} requires at least one condition image") + if type(batch_size) is not int or batch_size <= 0: + raise ValueError( + f"{type(adapter).__name__} condition batch_size must be a positive int, " + f"received {batch_size!r}" + ) + vae_dtype = getattr(adapter.vae, "dtype", dtype) + cond_tokens: List[torch.Tensor] = [] + cond_ids: List[torch.Tensor] = [] + for index, image in enumerate(images): + encoded = encode_flux2_vae_image( + adapter, + image.to(device=device, dtype=vae_dtype), + sample_mode="argmax", + ) + latent_height, latent_width = encoded.shape[-2:] + cond_tokens.append(pack_flux3_latents(encoded).squeeze(0)) + cond_ids.append( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=image_rope_scale, + l_coord=-(1 + index), + device=device, + ) + ) + packed = torch.cat(cond_tokens, dim=0).unsqueeze(0).repeat(batch_size, 1, 1) + packed = packed.to(device=device, dtype=dtype) + image_latent_ids = torch.cat(cond_ids, dim=0).unsqueeze(0).repeat(batch_size, 1, 1).to(device) + return packed, image_latent_ids + + __all__ = [ "encode_flux1_vae_image", "encode_flux2_output_images", "encode_flux2_vae_image", + "encode_flux3_output_images", + "pack_flux3_latents", + "patchify_flux2_latents", "prepare_flux1_output_latents", "prepare_flux2_condition_latents", + "prepare_flux3_condition_latents", + "prepare_flux3_image_ids", + "unpack_flux3_latents", + "unpatchify_flux2_latents", ] diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py new file mode 100644 index 00000000..5e99f9ad --- /dev/null +++ b/src/flow_factory/models/flux/flux3.py @@ -0,0 +1,1563 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# src/flow_factory/models/flux/flux3.py +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +from accelerate import Accelerator +from PIL import Image + +import diffusers +from diffusers.utils.torch_utils import randn_tensor + +try: + from diffusers import Flux3Transformer2DModel, ModularPipeline + from diffusers.modular_pipelines.flux3 import ( + Flux3AutoBlocks, + Flux3DistilledAutoBlocks, + Flux3DistilledModularPipeline, + Flux3ModularPipeline, + ) + from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor + + _FLUX3_IMPORT_ERROR: Optional[Exception] = None +except (ImportError, RuntimeError) as exc: # pragma: no cover - exercised with older diffusers + Flux3Transformer2DModel = Any # type: ignore[misc,assignment] + ModularPipeline = Any # type: ignore[misc,assignment] + Flux3AutoBlocks = Any # type: ignore[misc,assignment] + Flux3DistilledAutoBlocks = Any # type: ignore[misc,assignment] + Flux3DistilledModularPipeline = Any # type: ignore[misc,assignment] + Flux3ModularPipeline = Any # type: ignore[misc,assignment] + Flux3ImageProcessor = Any # type: ignore[misc,assignment] + _FLUX3_IMPORT_ERROR = exc + +from ...contracts import InputMediaOrder, NegativePromptPolicy, PipelineIOContract +from ...hparams import Arguments +from ...samples import I2ISample +from ...scheduler import ( + FlowMatchEulerDiscreteSDEScheduler, + FlowMatchEulerDiscreteSDESchedulerOutput, + set_scheduler_timesteps, +) +from ...utils.image import ( + ImageBatch, + ImageSingle, + MultiImageBatch, + is_multi_image_batch, + standardize_image_batch, +) +from ...utils.logger_utils import setup_logger +from ...utils.trajectory_collector import ( + TrajectoryIndicesType, + create_callback_collector, + create_trajectory_collector, +) +from ..abc import BaseAdapter +from ..checkpointing import CheckpointUnit +from ..configured_image_output import ( + ConfiguredImageOutputAdapterMixin, + EncodedImageTensor, +) +from ..pipeline_contracts import image_output_contract +from ..precision import DTypePolicy +from ..runtime import ModularPipelineRuntime +from ._output import ( + encode_flux3_output_images, + pack_flux3_latents, + prepare_flux3_condition_latents, + prepare_flux3_image_ids, + unpack_flux3_latents, + unpatchify_flux2_latents, +) + +logger = setup_logger(__name__) + +FLUX3_DIFFUSERS_COMMIT = "7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941" +FLUX3_INSTALL = "git submodule update --init && pip install -e ./diffusers" +CONDITION_IMAGE_SIZE = (1024, 1024) +IMAGE_ROPE_SCALE = 0.5 +TEXT_PAD_MULTIPLE = 80 +DEFAULT_TEXT_ENCODER_OUT_LAYERS: Tuple[int, ...] = (4, 8, 12, 16, 20, 24, 28, 32) + + +@dataclass +class Flux3Sample(I2ISample): + """Output sample class for Flux3Adapter models.""" + + _shared_fields: ClassVar[frozenset[str]] = frozenset({}) + latent_ids: Optional[torch.Tensor] = None + text_ids: Optional[torch.Tensor] = None + negative_text_ids: Optional[torch.Tensor] = None + image_latents: Optional[torch.Tensor] = None + image_latent_ids: Optional[torch.Tensor] = None + + +class Flux3Adapter(ConfiguredImageOutputAdapterMixin, BaseAdapter): + """Adapter for FLUX.3 modular text-to-image and multi-reference image-to-image models.""" + + supports_diffusers_cache: ClassVar[bool] = False + supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = True + ddp_find_unused_parameters: ClassVar[bool] = True + component_load_dtype_defaults: ClassVar[DTypePolicy] = MappingProxyType( + { + "transformers": torch.bfloat16, + "text_encoders": torch.bfloat16, + "vae": torch.float32, + } + ) + offline_training_forward_overrides = MappingProxyType({"guidance_scale": 1.0}) + pipeline_io_contract = image_output_contract( + negative_prompt=NegativePromptPolicy.OPTIONAL, + input_image_min_count=0, + input_image_max_count=None, + input_order=InputMediaOrder.WITHIN_TYPE, + ) + + def __init__(self, config: Arguments, accelerator: Accelerator): + if _FLUX3_IMPORT_ERROR is not None: + raise ImportError( + "Flux3Adapter requires the bundled Diffusers submodule at " + f"{FLUX3_DIFFUSERS_COMMIT}. Install it with `{FLUX3_INSTALL}`. " + f"The active diffusers version is {getattr(diffusers, '__version__', 'unknown')}." + ) from _FLUX3_IMPORT_ERROR + super().__init__(config, accelerator) + self.pipeline: Flux3ModularPipeline + self.scheduler: FlowMatchEulerDiscreteSDEScheduler + self._has_warned_inference_fallback = False + self._has_warned_forward_fallback = False + + def load_pipeline(self) -> Flux3ModularPipeline: + """Load the FLUX.3 modular pipeline specification from a local path or Hub repo.""" + if _FLUX3_IMPORT_ERROR is not None: + raise ImportError( + "Flux3Adapter requires the bundled Diffusers submodule at " + f"{FLUX3_DIFFUSERS_COMMIT}. Install it with `{FLUX3_INSTALL}`. " + f"The active diffusers version is {getattr(diffusers, '__version__', 'unknown')}." + ) from _FLUX3_IMPORT_ERROR + model_name_or_path = self.model_args.model_name_or_path + pipeline = ModularPipeline.from_pretrained(model_name_or_path) + if Path(model_name_or_path).is_dir(): + for name in pipeline.pretrained_component_names: + spec = pipeline._component_specs[name] + spec.pretrained_model_name_or_path = model_name_or_path + spec.revision = None + return pipeline + + def build_component_runtime(self) -> ModularPipelineRuntime: + """Build the modular runtime and materialize the target transformer.""" + runtime = ModularPipelineRuntime.from_adapter(self, self.load_pipeline()) + target_names = runtime.resolve_component_names(self.model_args.target_components) + runtime.materialize_components(target_names) + return runtime + + @property + def is_distilled(self) -> bool: + """Return whether the loaded checkpoint uses the step-distilled flow-map contract.""" + transformer = ( + self.get_component("transformer") if self.has_component("transformer") else None + ) + transformer_config = getattr(transformer, "config", None) + if getattr(transformer_config, "flow_map", None) is not None: + return bool(transformer_config.flow_map) + if Flux3DistilledModularPipeline is not Any and isinstance( + self.pipeline, Flux3DistilledModularPipeline + ): + return True + return "guider" not in self.component_runtime.declared_component_names + + def _resolve_pipeline_io_contract(self) -> Optional[PipelineIOContract]: + """Narrow negative-prompt support when a step-distilled flow-map checkpoint is loaded.""" + if self.is_distilled: + return image_output_contract( + negative_prompt=NegativePromptPolicy.UNSUPPORTED, + input_image_min_count=0, + input_image_max_count=None, + input_order=InputMediaOrder.WITHIN_TYPE, + ) + return type(self).pipeline_io_contract + + @property + def preprocessing_modules(self) -> List[str]: + return ["text_encoder", "tokenizer", "vae", "image_processor"] + + @property + def inference_modules(self) -> List[str]: + return ["transformer", "vae", "image_processor"] + + @property + def tokenizer(self) -> Any: + """Expose the processor-owned tokenizer to distributed reward decoding.""" + if self.get_component("tokenizer") is None and self.has_component("tokenizer"): + self.component_runtime.materialize_components(["tokenizer"]) + processor = self.get_component("tokenizer") + if processor is None: + raise ValueError("No tokenizer found in the FLUX.3 pipeline.") + return getattr(processor, "tokenizer", processor) + + @property + def processor(self) -> Any: + """Return the Qwen3-VL processor declared under the modular ``tokenizer`` slot.""" + if self.get_component("tokenizer") is None and self.has_component("tokenizer"): + self.component_runtime.materialize_components(["tokenizer"]) + processor = self.get_component("tokenizer") + if processor is None: + raise ValueError("No processor/tokenizer component found in the FLUX.3 pipeline.") + return processor + + @property + def default_target_modules(self) -> List[str]: + """Default LoRA target modules across all FLUX.3 mode and single-stream blocks.""" + return [ + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.linear_in", + "ff.linear_out", + ] + + def _gradient_checkpointing_units( + self, + component_name: str, + component: torch.nn.Module, + ) -> List[CheckpointUnit]: + """Return FLUX.3 transformer blocks in forward execution order.""" + txt_blocks = getattr(component, "txt_mode_blocks", None) + mode_blocks = getattr(component, "mode_blocks", None) + single_blocks = getattr(component, "single_transformer_blocks", None) + image_blocks = ( + mode_blocks["image"] + if isinstance(mode_blocks, nn.ModuleDict) and "image" in mode_blocks + else None + ) + cond_blocks = ( + mode_blocks["image_cond"] + if isinstance(mode_blocks, nn.ModuleDict) and "image_cond" in mode_blocks + else None + ) + if ( + isinstance(txt_blocks, nn.ModuleList) + and isinstance(image_blocks, nn.ModuleList) + and isinstance(cond_blocks, nn.ModuleList) + and isinstance(single_blocks, nn.ModuleList) + ): + return [ + *[(f"txt_mode_blocks.{idx}", block) for idx, block in enumerate(txt_blocks)], + *[(f"mode_blocks.image.{idx}", block) for idx, block in enumerate(image_blocks)], + *[ + (f"mode_blocks.image_cond.{idx}", block) + for idx, block in enumerate(cond_blocks) + ], + *[ + (f"single_transformer_blocks.{idx}", block) + for idx, block in enumerate(single_blocks) + ], + ] + return super()._gradient_checkpointing_units(component_name, component) + + # ======================== Prompt Encoding ======================== + + @staticmethod + def _prepare_text_ids(seq_len: int, device: Optional[torch.device] = None) -> torch.Tensor: + """Prepare 4D position IDs ``(T, H, W, L)`` of shape ``(seq_len, 4)`` for text tokens.""" + if seq_len <= 0: + raise ValueError(f"expected positive text sequence length, received {seq_len}") + ids = torch.cartesian_prod( + torch.arange(1), + torch.arange(1), + torch.arange(1), + torch.arange(seq_len), + ).to(dtype=torch.float32) + if device is not None: + ids = ids.to(device) + return ids + + @staticmethod + def _encode_single_qwen3_vl_prompt( + text_encoder: Any, + processor: Any, + prompt: str, + *, + dtype: torch.dtype, + device: torch.device, + max_sequence_length: int = 512, + pad_multiple: int = TEXT_PAD_MULTIPLE, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Encode one prompt with the FLUX.3 80-token right-padded Qwen3-VL recipe.""" + if not isinstance(prompt, str): + raise TypeError( + f"expected prompt to be str, received {type(prompt).__name__}: {prompt!r}" + ) + if max_sequence_length <= 0: + raise ValueError( + f"expected positive max_sequence_length, received {max_sequence_length}" + ) + if pad_multiple <= 0: + raise ValueError(f"expected positive pad_multiple, received {pad_multiple}") + if not text_encoder_out_layers: + raise ValueError("expected non-empty text_encoder_out_layers") + + inner_tokenizer = getattr(processor, "tokenizer", processor) + chat_text = processor.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + ) + unpadded_ids = inner_tokenizer( + chat_text, + return_tensors="pt", + truncation=True, + max_length=max_sequence_length, + )["input_ids"] + real_length = max(int(unpadded_ids.shape[1]), 1) + padded_length = min( + math.ceil(real_length / pad_multiple) * pad_multiple, + max_sequence_length, + ) + inputs = inner_tokenizer( + chat_text, + return_tensors="pt", + padding="max_length", + padding_side="right", + truncation=True, + max_length=padded_length, + ) + + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + # Qwen3-VL computes padding position IDs as 1 in transformers 4.57 (the FLUX.3 training + # environment) and 0 in 5.x. Pass explicit position IDs so padding tokens match upstream. + position_ids = torch.arange(padded_length, device=device)[None].masked_fill( + attention_mask == 0, 1 + ) + + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=True, + use_cache=False, + ) + + stacked = torch.stack([output.hidden_states[k] for k in text_encoder_out_layers], dim=1) + stacked = stacked.to(dtype=dtype, device=device) + batch_size, num_layers, seq_len, hidden_dim = stacked.shape + prompt_embeds = stacked.permute(0, 2, 1, 3).reshape( + batch_size, seq_len, num_layers * hidden_dim + ) + return input_ids.squeeze(0), prompt_embeds.squeeze(0) + + @classmethod + def _get_qwen3_vl_prompt_embeds( + cls, + text_encoder: Any, + processor: Any, + prompt: Union[str, Sequence[str]], + *, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + max_sequence_length: int = 512, + pad_multiple: int = TEXT_PAD_MULTIPLE, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + ) -> Tuple[ + Union[torch.Tensor, List[torch.Tensor]], + Union[torch.Tensor, List[torch.Tensor]], + Union[torch.Tensor, List[torch.Tensor]], + ]: + """Encode one or more prompts while preserving each prompt's exact 80-token window.""" + resolved_dtype = text_encoder.dtype if dtype is None else dtype + resolved_device = text_encoder.device if device is None else device + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + if not prompts: + raise ValueError("expected at least one prompt to encode") + + ids_list: List[torch.Tensor] = [] + embeds_list: List[torch.Tensor] = [] + text_ids_list: List[torch.Tensor] = [] + for single_prompt in prompts: + single_ids, single_embeds = cls._encode_single_qwen3_vl_prompt( + text_encoder=text_encoder, + processor=processor, + prompt=single_prompt, + dtype=resolved_dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + ids_list.append(single_ids) + embeds_list.append(single_embeds) + text_ids_list.append( + cls._prepare_text_ids(single_embeds.shape[0], device=resolved_device) + ) + + lengths = {item.shape[0] for item in embeds_list} + if len(lengths) == 1: + return ( + torch.stack(ids_list, dim=0), + torch.stack(embeds_list, dim=0), + torch.stack(text_ids_list, dim=0), + ) + return ids_list, embeds_list, text_ids_list + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + guidance_scale: float = 3.5, + device: Optional[torch.device] = None, + max_sequence_length: int = 512, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]: + """Encode prompt(s) and optional negative prompt(s) with Qwen3-VL.""" + if ( + self.get_component("text_encoder") is None or self.get_component("tokenizer") is None + ) and self.has_component("text_encoder"): + self.on_load_components(["text_encoder", "tokenizer"], device=device) + + text_encoder = self._require_component("text_encoder") + processor = self.processor + resolved_device = text_encoder.device if device is None else device + + if prompt is None: + prompt = "" + prompt_list = [prompt] if isinstance(prompt, str) else list(prompt) + + if self.is_distilled and negative_prompt is not None: + neg_values = ( + [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt) + ) + if any(item is not None and item != "" for item in neg_values): + raise ValueError( + "Step-distilled FLUX.3 (flow_map=True) does not support negative_prompt; " + f"received {negative_prompt!r}." + ) + + do_classifier_free_guidance = (not self.is_distilled) and guidance_scale > 1.0 + + prompt_ids, prompt_embeds, text_ids = self._get_qwen3_vl_prompt_embeds( + text_encoder=text_encoder, + processor=processor, + prompt=prompt_list, + dtype=text_encoder.dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + results: Dict[str, Union[torch.Tensor, List[torch.Tensor]]] = { + "prompt_ids": prompt_ids, + "prompt_embeds": prompt_embeds, + "text_ids": text_ids, + } + + if do_classifier_free_guidance: + negative_prompt = "" if negative_prompt is None else negative_prompt + neg_list = ( + [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt) + ) + if len(neg_list) == 1 and len(prompt_list) > 1: + neg_list = neg_list * len(prompt_list) + if len(neg_list) != len(prompt_list): + raise ValueError( + f"The number of negative prompts ({len(neg_list)}) must match " + f"the number of prompts ({len(prompt_list)})." + ) + neg_ids, neg_embeds, neg_text_ids = self._get_qwen3_vl_prompt_embeds( + text_encoder=text_encoder, + processor=processor, + prompt=neg_list, + dtype=text_encoder.dtype, + device=resolved_device, + max_sequence_length=max_sequence_length, + pad_multiple=pad_multiple, + text_encoder_out_layers=text_encoder_out_layers, + ) + results.update( + { + "negative_prompt_ids": neg_ids, + "negative_prompt_embeds": neg_embeds, + "negative_text_ids": neg_text_ids, + } + ) + + return results + + # ======================== Image Encoding ======================== + + def encode_image( + self, + images: Union[ImageSingle, ImageBatch, MultiImageBatch], + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + generator: Optional[torch.Generator] = None, + ) -> Dict[str, Union[List[List[torch.Tensor]], List[Optional[torch.Tensor]]]]: + """Encode reference image(s) into normalized packed FLUX.3 conditioning tokens.""" + del generator + if self.get_component("vae") is None and self.has_component("vae"): + self.on_load_components(["vae", "image_processor"], device=device) + + vae = self._require_component("vae") + resolved_device = vae.device if device is None else device + transformer = ( + self.get_component("transformer") if self.has_component("transformer") else None + ) + resolved_dtype = ( + dtype if dtype is not None else getattr(transformer, "dtype", None) or vae.dtype + ) + + if not self._is_multi_images_batch(images): + images = [images] # type: ignore[list-item] + + standardized_batches = [ + ( + [] + if isinstance(imgs, list) and not imgs + else self._standardize_image_input(imgs, output_type="pil") + ) + for imgs in images # type: ignore[union-attr] + ] + + condition_image_tensors: List[List[torch.Tensor]] = [ + ( + self._resize_condition_images( + condition_images=imgs, + condition_image_size=condition_image_size, + ) + if imgs + else [] + ) + for imgs in standardized_batches + ] + + image_latents_list: List[Optional[torch.Tensor]] = [] + image_latent_ids_list: List[Optional[torch.Tensor]] = [] + for cond_img_tensors in condition_image_tensors: + if not cond_img_tensors: + image_latents_list.append(None) + image_latent_ids_list.append(None) + continue + image_latents, image_latent_ids = prepare_flux3_condition_latents( + self, + cond_img_tensors, + batch_size=1, + device=resolved_device, + dtype=resolved_dtype, + image_rope_scale=IMAGE_ROPE_SCALE, + ) + image_latents_list.append(image_latents.squeeze(0)) + image_latent_ids_list.append(image_latent_ids.squeeze(0)) + + postprocessed_condition_images: List[List[torch.Tensor]] = [ + [ + self.pipeline.image_processor.postprocess(img, output_type="pt")[0].to( + resolved_device + ) + for img in cond_img_tensors + ] + for cond_img_tensors in condition_image_tensors + ] + + return { + "condition_images": postprocessed_condition_images, + "image_latents": image_latents_list, + "image_latent_ids": image_latent_ids_list, + } + + def _resize_condition_images( + self, + condition_images: Union[Image.Image, List[Image.Image]], + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + ) -> List[torch.Tensor]: + """Resize reference images when exceeding the area budget and center-crop to 16px grid.""" + if isinstance(condition_images, Image.Image): + condition_images = [condition_images] + + condition_images = self._standardize_image_input( + condition_images, + output_type="pil", + ) + + if isinstance(condition_image_size, int): + condition_image_size = (condition_image_size, condition_image_size) + elif not isinstance(condition_image_size, (tuple, list)) or len(condition_image_size) != 2: + raise ValueError( + f"expected condition_image_size as int or (height, width), " + f"received {condition_image_size!r}" + ) + max_area = int(condition_image_size[0]) * int(condition_image_size[1]) + if max_area <= 0: + raise ValueError( + f"expected positive condition_image_size area, received {condition_image_size!r}" + ) + + multiple_of = self.pipeline.vae_scale_factor * 2 + condition_image_tensors: List[torch.Tensor] = [] + for img in condition_images: + image_width, image_height = img.size + if image_width * image_height > max_area: + img = self.pipeline.image_processor._resize_to_target_area(img, max_area) + image_width, image_height = img.size + + cropped_width = max(multiple_of, (image_width // multiple_of) * multiple_of) + cropped_height = max(multiple_of, (image_height // multiple_of) * multiple_of) + tensor_img = self.pipeline.image_processor.preprocess( + img, + height=cropped_height, + width=cropped_width, + resize_mode="crop", + ) + condition_image_tensors.append(tensor_img) + + return condition_image_tensors + + @staticmethod + def _is_multi_images_batch(images: Optional[Union[ImageBatch, MultiImageBatch]]) -> bool: + return images is not None and is_multi_image_batch(images) + + @staticmethod + def _is_multi_image_latents( + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]], + ) -> bool: + return ( + isinstance(image_latents, list) + and len(image_latents) > 0 + and all( + latent is None or (isinstance(latent, torch.Tensor) and latent.ndim == 2) + for latent in image_latents + ) + ) or (isinstance(image_latents, torch.Tensor) and image_latents.ndim == 3) + + def _standardize_image_input( + self, + images: Union[ImageSingle, ImageBatch], + output_type: Literal["pil", "pt", "np"] = "pil", + ) -> ImageBatch: + """Standardize image input to the requested container format.""" + if isinstance(images, Image.Image): + images = [images] + return standardize_image_batch(images, output_type=output_type) + + # ======================== Output Image Encoding & Decoding ======================== + + def _output_geometry_multiple(self) -> int: + """Require the 16-pixel grid imposed by the 8x VAE and 2x2 latent patchification.""" + return self.pipeline.vae_scale_factor * 2 + + def _encode_output_images( + self, + pixel_values: torch.Tensor, + condition: Mapping[str, Any], + generator: Optional[torch.Generator], + ) -> EncodedImageTensor: + """Sample and pack target images with the FLUX.3 latent recipe.""" + del condition + return encode_flux3_output_images( + self, + pixel_values, + generator, + image_rope_scale=IMAGE_ROPE_SCALE, + ) + + def _resolve_latent_grid( + self, + latents: torch.Tensor, + *, + latent_ids: Optional[torch.Tensor] = None, + height: Optional[int] = None, + width: Optional[int] = None, + ) -> Tuple[int, int]: + """Resolve ``(latent_height, latent_width)`` from explicit dimensions or position IDs.""" + scale = self.pipeline.vae_scale_factor * 2 + seq_len = latents.shape[1] + if height is not None and width is not None: + latent_height = int(height) // scale + latent_width = int(width) // scale + if latent_height * latent_width != seq_len: + raise ValueError( + f"FLUX.3 latents has {seq_len} tokens, which does not match " + f"height x width = {height} x {width} ({latent_height * latent_width} tokens)" + ) + return latent_height, latent_width + + if latent_ids is not None: + ids_2d = self._as_2d_position_ids(latent_ids, "latent_ids") + max_h = float(ids_2d[:, 1].max().item()) + max_w = float(ids_2d[:, 2].max().item()) + latent_height = int(round(max_h / IMAGE_ROPE_SCALE)) + 1 + latent_width = int(round(max_w / IMAGE_ROPE_SCALE)) + 1 + if latent_height * latent_width == seq_len: + return latent_height, latent_width + + side = int(math.isqrt(seq_len)) + if side * side == seq_len: + return side, side + raise ValueError( + f"Cannot infer FLUX.3 2D latent grid for seq_len={seq_len}; " + "pass explicit height/width or valid latent_ids." + ) + + def decode_latents( + self, + latents: torch.Tensor, + latent_ids: Optional[torch.Tensor] = None, + height: Optional[int] = None, + width: Optional[int] = None, + output_type: Literal["pil", "pt", "np"] = "pil", + ) -> Union[List[Image.Image], torch.Tensor, np.ndarray]: + """Unpack, BatchNorm-denormalize, unpatchify, and decode FLUX.3 image latents.""" + if self.get_component("vae") is None and self.has_component("vae"): + self.on_load_components(["vae", "image_processor"], device=latents.device) + + vae = self._require_component("vae") + latent_height, latent_width = self._resolve_latent_grid( + latents, + latent_ids=latent_ids, + height=height, + width=width, + ) + unpacked = unpack_flux3_latents(latents, latent_height, latent_width) + + latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(unpacked.device, unpacked.dtype) + latents_bn_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps + ).to(unpacked.device, unpacked.dtype) + unpacked = unpacked * latents_bn_std + latents_bn_mean + unpatched = unpatchify_flux2_latents(unpacked) + + decoded = vae.decode(unpatched.to(device=vae.device, dtype=vae.dtype), return_dict=False)[0] + return self.pipeline.image_processor.postprocess(decoded, output_type=output_type) + + # ======================== Inference ======================== + + @staticmethod + def _prepare_latents( + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + *, + vae_scale_factor: int, + dtype: torch.dtype, + device: torch.device, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Draw initial float32 noise on the 2D grid, cast, pack to tokens, and build RoPE IDs.""" + image_scale_factor = vae_scale_factor * 2 + if height % image_scale_factor != 0 or width % image_scale_factor != 0: + raise ValueError( + f"`height` and `width` must be divisible by {image_scale_factor}, " + f"received {height} and {width}." + ) + latent_height = height // image_scale_factor + latent_width = width // image_scale_factor + if latents is None: + shape = (batch_size, num_channels_latents, latent_height, latent_width) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"Passed a list of {len(generator)} generators for batch_size={batch_size}." + ) + latents = randn_tensor( + shape, + generator=generator, + device=device, + dtype=torch.float32, + ).to(dtype) + else: + latents = latents.to(device=device, dtype=dtype) + if latents.ndim == 4: + latents = pack_flux3_latents(latents) + latent_ids = ( + prepare_flux3_image_ids( + latent_height, + latent_width, + image_rope_scale=IMAGE_ROPE_SCALE, + l_coord=0, + device=device, + ) + .unsqueeze(0) + .repeat(batch_size, 1, 1) + ) + return latents, latent_ids + + @torch.no_grad() + def _inference( + self, + images: Optional[Union[ImageBatch, MultiImageBatch]] = None, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_ids: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + text_ids: Optional[torch.Tensor] = None, + negative_prompt_ids: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_text_ids: Optional[torch.Tensor] = None, + condition_images: Optional[MultiImageBatch] = None, + image_latents: Optional[torch.Tensor] = None, + image_latent_ids: Optional[torch.Tensor] = None, + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + max_sequence_length: int = 512, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + compute_log_prob: bool = False, + extra_call_back_kwargs: List[str] = [], + trajectory_indices: TrajectoryIndicesType = "all", + ) -> List[Flux3Sample]: + """Run FLUX.3 rollout on a uniform-window T2I batch or a single sample.""" + device = self.device + transformer = self._require_component("transformer") + dtype = transformer.dtype + do_classifier_free_guidance = (not self.is_distilled) and guidance_scale > 1.0 + resolved_attention_kwargs = ( + attention_kwargs if attention_kwargs is not None else joint_attention_kwargs + ) + + # 1. Encode prompt + if isinstance(prompt, str): + prompt = [prompt] + + if (prompt_embeds is None or prompt_ids is None or text_ids is None) or ( + do_classifier_free_guidance + and ( + negative_prompt_embeds is None + or negative_prompt_ids is None + or negative_text_ids is None + ) + ): + prompt_encoding = self.encode_prompt( + prompt=prompt or [""], + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + device=device, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + ) + prompt_ids = prompt_encoding["prompt_ids"] # type: ignore[assignment] + prompt_embeds = prompt_encoding["prompt_embeds"] # type: ignore[assignment] + text_ids = prompt_encoding["text_ids"] # type: ignore[assignment] + if do_classifier_free_guidance: + negative_prompt_ids = prompt_encoding["negative_prompt_ids"] # type: ignore[assignment] + negative_prompt_embeds = prompt_encoding["negative_prompt_embeds"] # type: ignore[assignment] + negative_text_ids = prompt_encoding["negative_text_ids"] # type: ignore[assignment] + else: + prompt_ids = prompt_ids.to(device) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + text_ids = text_ids.to(device) + if negative_prompt_ids is not None: + negative_prompt_ids = negative_prompt_ids.to(device) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + if negative_text_ids is not None: + negative_text_ids = negative_text_ids.to(device) + + if ( + not isinstance(prompt_embeds, torch.Tensor) + or not isinstance(prompt_ids, torch.Tensor) + or not isinstance(text_ids, torch.Tensor) + ): + raise ValueError( + "Flux3Adapter._inference requires uniform stacked prompt tensors; " + "use Flux3Adapter.inference() for variable-window batches." + ) + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + if isinstance(negative_prompt_embeds, torch.Tensor): + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + + batch_size = prompt_embeds.shape[0] + + # 2. Encode condition images + images = ( + [images] if images is not None and not self._is_multi_images_batch(images) else images + ) + if images is not None and ( + condition_images is None or image_latents is None or image_latent_ids is None + ): + image_encoding = self.encode_image( + images=images, + condition_image_size=condition_image_size, + device=device, + dtype=dtype, + generator=generator[0] if isinstance(generator, list) else generator, + ) + condition_images = image_encoding["condition_images"] # type: ignore[assignment] + raw_latents = image_encoding["image_latents"] + raw_ids = image_encoding["image_latent_ids"] + if any(item is None for item in raw_latents): + image_latents = None + image_latent_ids = None + else: + image_latents = torch.stack( + [item for item in raw_latents if item is not None], dim=0 + ).to(device=device, dtype=dtype) + image_latent_ids = torch.stack( + [item for item in raw_ids if item is not None], dim=0 + ).to(device=device) + else: + image_latents = ( + image_latents.to(device=device, dtype=dtype) if image_latents is not None else None + ) + image_latent_ids = ( + image_latent_ids.to(device=device) if image_latent_ids is not None else None + ) + + # 3. Prepare initial latents + num_channels_latents = int(getattr(transformer.config, "in_channels", 128)) + latents, latent_ids = self._prepare_latents( + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + vae_scale_factor=self.pipeline.vae_scale_factor, + dtype=dtype, + device=device, + generator=generator, + ) + + # 4. Prepare timesteps (constant shift 6.93 on the scheduler) + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + timesteps = set_scheduler_timesteps( + scheduler=self.scheduler, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + seq_len=latents.shape[1], + device=device, + ) + + # 5. Denoising loop + latent_collector = create_trajectory_collector(trajectory_indices, num_inference_steps) + latents = self.cast_latents(latents, default_dtype=dtype) + latent_collector.collect(latents, step_idx=0) + if compute_log_prob: + log_prob_collector = create_trajectory_collector( + trajectory_indices, num_inference_steps + ) + callback_collector = create_callback_collector(trajectory_indices, num_inference_steps) + + for i, t in enumerate(timesteps): + current_noise_level = self.scheduler.get_noise_level_for_timestep(t) + t_next = ( + timesteps[i + 1] + if i + 1 < len(timesteps) + else torch.tensor(0, device=device, dtype=timesteps.dtype) + ) + return_kwargs = list( + set(["next_latents", "log_prob", "velocity"] + extra_call_back_kwargs) + ) + current_compute_log_prob = compute_log_prob and current_noise_level > 0 + + output = self._forward( + t=t, + t_next=t_next, + latents=latents, + latent_ids=latent_ids, + prompt_embeds=prompt_embeds, + text_ids=text_ids, + image_latents=image_latents, + image_latent_ids=image_latent_ids, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + guidance_scale=guidance_scale, + attention_kwargs=resolved_attention_kwargs, + compute_log_prob=current_compute_log_prob, + return_kwargs=return_kwargs, + noise_level=current_noise_level, + ) + + latents = self.cast_latents(output.next_latents, default_dtype=dtype) + latent_collector.collect(latents, i + 1) + if current_compute_log_prob: + log_prob_collector.collect(output.log_prob, i) + + callback_collector.collect_step( + step_idx=i, + output=output, + keys=extra_call_back_kwargs, + capturable={"noise_level": current_noise_level}, + ) + + # 6. Decode latents to images + decoded_images = self.decode_latents( + latents, + latent_ids=latent_ids, + height=height, + width=width, + output_type="pt", + ) + + # 7. Assemble Flux3Sample instances + extra_call_back_res = callback_collector.get_result() + callback_index_map = callback_collector.get_index_map() + all_latents = latent_collector.get_result() + latent_index_map = latent_collector.get_index_map() + all_log_probs = log_prob_collector.get_result() if compute_log_prob else None + log_prob_index_map = log_prob_collector.get_index_map() if compute_log_prob else None + + samples = [ + Flux3Sample( + timesteps=timesteps, + all_latents=( + torch.stack([lat[b] for lat in all_latents], dim=0) + if all_latents is not None + else None + ), + log_probs=( + torch.stack([lp[b] for lp in all_log_probs], dim=0) + if all_log_probs is not None + else None + ), + latent_index_map=latent_index_map, + log_prob_index_map=log_prob_index_map, + height=height, + width=width, + image=decoded_images[b], + latent_ids=latent_ids[b], + prompt=prompt[b] if isinstance(prompt, list) else prompt, + prompt_ids=prompt_ids[b], + prompt_embeds=prompt_embeds[b], + text_ids=text_ids[b], + negative_prompt=( + negative_prompt[b] if isinstance(negative_prompt, list) else negative_prompt + ), + negative_prompt_ids=( + negative_prompt_ids[b] if negative_prompt_ids is not None else None + ), + negative_prompt_embeds=( + negative_prompt_embeds[b] if negative_prompt_embeds is not None else None + ), + negative_text_ids=(negative_text_ids[b] if negative_text_ids is not None else None), + condition_images=condition_images[b] if condition_images is not None else None, + image_latents=image_latents[b] if image_latents is not None else None, + image_latent_ids=image_latent_ids[b] if image_latent_ids is not None else None, + extra_kwargs={ + **{k: v[b] for k, v in extra_call_back_res.items()}, + "callback_index_map": callback_index_map, + }, + ) + for b in range(batch_size) + ] + + if hasattr(self.pipeline, "maybe_free_model_hooks"): + self.pipeline.maybe_free_model_hooks() + + return samples + + @torch.no_grad() + def inference( + self, + images: Optional[MultiImageBatch] = None, + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 3.5, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + max_sequence_length: int = 512, + text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, + pad_multiple: int = TEXT_PAD_MULTIPLE, + prompt_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + prompt_embeds: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + text_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_prompt_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_prompt_embeds: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + negative_text_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, + condition_images: Optional[MultiImageBatch] = None, + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + image_latent_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + compute_log_prob: bool = False, + extra_call_back_kwargs: List[str] = [], + trajectory_indices: TrajectoryIndicesType = "all", + ) -> List[Flux3Sample]: + """Run FLUX.3 inference; falls back to per-sample execution for I2I or ragged windows.""" + if isinstance(prompt, str): + prompt = [prompt] + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * (len(prompt) if prompt is not None else 1) + + is_nested_images = self._is_multi_images_batch(images) + is_nested_image_latents = self._is_multi_image_latents(image_latents) + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + # If raw prompts are passed with batch_size > 1 and no pre-encoded embeds, check whether + # their 80-token windows differ by encoding once upfront. + if ( + prompt is not None + and len(prompt) > 1 + and prompt_embeds is None + and not (is_nested_images or is_nested_image_latents) + ): + encoded = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + device=self.device, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + ) + prompt_ids = encoded["prompt_ids"] + prompt_embeds = encoded["prompt_embeds"] + text_ids = encoded["text_ids"] + negative_prompt_ids = encoded.get("negative_prompt_ids") + negative_prompt_embeds = encoded.get("negative_prompt_embeds") + negative_text_ids = encoded.get("negative_text_ids") + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + fall_back = is_nested_images or is_nested_image_latents or is_ragged_prompts + if not fall_back: + return self._inference( + images=images, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + condition_image_size=condition_image_size, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + prompt_ids=prompt_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + negative_prompt_ids=negative_prompt_ids, # type: ignore[arg-type] + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + condition_images=condition_images, + image_latents=image_latents, # type: ignore[arg-type] + image_latent_ids=image_latent_ids, # type: ignore[arg-type] + compute_log_prob=compute_log_prob, + extra_call_back_kwargs=extra_call_back_kwargs, + trajectory_indices=trajectory_indices, + ) + + if not self._has_warned_inference_fallback: + logger.warning( + "FLUX.3 per-sample inference fallback activated (image conditioning or " + "variable 80-token prompt windows). This warning will only appear once." + ) + self._has_warned_inference_fallback = True + + if images is not None: + batch_size = len(images) + elif image_latents is not None: + batch_size = len(image_latents) + elif prompt_embeds is not None: + batch_size = len(prompt_embeds) + elif prompt is not None: + batch_size = len(prompt) + else: + raise ValueError("Flux3Adapter.inference requires prompt, prompt_embeds, or images.") + + samples: List[Flux3Sample] = [] + for idx in range(batch_size): + this_prompt = [prompt[idx]] if prompt is not None else None + this_neg_prompt = ( + [negative_prompt[idx]] if isinstance(negative_prompt, list) else negative_prompt + ) + this_prompt_ids = ( + prompt_ids[idx].unsqueeze(0) + if prompt_ids is not None and prompt_ids[idx] is not None + else None + ) + this_prompt_embeds = ( + prompt_embeds[idx].unsqueeze(0) + if prompt_embeds is not None and prompt_embeds[idx] is not None + else None + ) + this_text_ids = ( + text_ids[idx].unsqueeze(0) + if text_ids is not None and text_ids[idx] is not None + else None + ) + this_neg_prompt_ids = ( + negative_prompt_ids[idx].unsqueeze(0) + if negative_prompt_ids is not None and negative_prompt_ids[idx] is not None + else None + ) + this_neg_prompt_embeds = ( + negative_prompt_embeds[idx].unsqueeze(0) + if negative_prompt_embeds is not None and negative_prompt_embeds[idx] is not None + else None + ) + this_neg_text_ids = ( + negative_text_ids[idx].unsqueeze(0) + if negative_text_ids is not None and negative_text_ids[idx] is not None + else None + ) + this_images = images[idx] if images is not None else None + this_condition_images = ( + condition_images[idx : idx + 1] if condition_images is not None else None + ) + this_image_latents = ( + image_latents[idx].unsqueeze(0) + if image_latents is not None and image_latents[idx] is not None + else None + ) + this_image_latent_ids = ( + image_latent_ids[idx].unsqueeze(0) + if image_latent_ids is not None and image_latent_ids[idx] is not None + else None + ) + sample = self._inference( + images=this_images, + prompt=this_prompt, + negative_prompt=this_neg_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator[idx] if isinstance(generator, list) else generator, + condition_image_size=condition_image_size, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + prompt_ids=this_prompt_ids, + prompt_embeds=this_prompt_embeds, + text_ids=this_text_ids, + negative_prompt_ids=this_neg_prompt_ids, + negative_prompt_embeds=this_neg_prompt_embeds, + negative_text_ids=this_neg_text_ids, + condition_images=this_condition_images, + image_latents=this_image_latents, + image_latent_ids=this_image_latent_ids, + compute_log_prob=compute_log_prob, + extra_call_back_kwargs=extra_call_back_kwargs, + trajectory_indices=trajectory_indices, + ) + samples.extend(sample) + + return samples + + # ======================== Forward ======================== + + @staticmethod + def _as_2d_position_ids(ids: torch.Tensor, name: str) -> torch.Tensor: + """Normalize ``(seq_len, 4)`` or batched ``(B, seq_len, 4)`` position IDs to 2D.""" + if not isinstance(ids, torch.Tensor): + raise TypeError(f"expected {name} to be torch.Tensor, received {type(ids).__name__}") + if ids.ndim == 2 and ids.shape[-1] == 4: + return ids + if ids.ndim == 3 and ids.shape[-1] == 4: + return ids[0] + raise ValueError( + f"expected {name} to have shape (seq_len, 4) or (B, seq_len, 4), " + f"received {tuple(ids.shape)}" + ) + + def _forward( + self, + t: torch.Tensor, + latents: torch.Tensor, + latent_ids: torch.Tensor, + prompt_embeds: torch.Tensor, + text_ids: torch.Tensor, + image_latents: Optional[torch.Tensor] = None, + image_latent_ids: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_text_ids: Optional[torch.Tensor] = None, + guidance_scale: float = 3.5, + t_next: Optional[torch.Tensor] = None, + next_latents: Optional[torch.Tensor] = None, + noise_level: Optional[float] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + compute_log_prob: bool = True, + return_kwargs: Sequence[str] = ( + "velocity", + "next_latents", + "next_latents_mean", + "std_dev_t", + "dt", + "log_prob", + ), + ) -> FlowMatchEulerDiscreteSDESchedulerOutput: + """Core single-step FLUX.3 forward handling T2I, I2I, CFG, and step-distilled flow maps.""" + batch_size = latents.shape[0] + transformer = self.transformer + transformer_dtype = getattr(transformer, "dtype", latents.dtype) + resolved_attention_kwargs = ( + attention_kwargs if attention_kwargs is not None else joint_attention_kwargs + ) + + is_distilled = self.is_distilled + if not is_distilled and guidance_scale > 1.0 and negative_prompt_embeds is None: + logger.warning( + "Passed `guidance_scale=%s` > 1.0 to Flux3Adapter, but `negative_prompt_embeds` " + "is None. Falling back to no classifier-free guidance.", + guidance_scale, + ) + do_classifier_free_guidance = ( + not is_distilled and guidance_scale > 1.0 and negative_prompt_embeds is not None + ) + + # 1. Prepare inputs in transformer compute dtype and 2D position IDs + latent_model_input = latents.to(dtype=transformer_dtype) + cond_prompt_embeds = prompt_embeds.to(dtype=transformer_dtype) + img_ids_2d = self._as_2d_position_ids(latent_ids, "latent_ids").to(latents.device) + txt_ids_2d = self._as_2d_position_ids(text_ids, "text_ids").to(latents.device) + + image_cond_hidden_states: Optional[torch.Tensor] = None + image_cond_ids_2d: Optional[torch.Tensor] = None + if image_latents is not None: + if image_latent_ids is None: + raise ValueError("`image_latents` requires `image_latent_ids` in Flux3Adapter.") + image_cond_hidden_states = image_latents.to( + device=latents.device, dtype=transformer_dtype + ) + image_cond_ids_2d = self._as_2d_position_ids(image_latent_ids, "image_latent_ids").to( + latents.device + ) + + timestep = (t.expand(batch_size) / 1000.0).to( + device=latents.device, dtype=transformer_dtype + ) + target_timestep: Optional[torch.Tensor] = None + if is_distilled: + target_t = ( + torch.zeros(batch_size, device=latents.device, dtype=t.dtype) + if t_next is None + else t_next.expand(batch_size) + ) + target_timestep = (target_t / 1000.0).to(device=latents.device, dtype=transformer_dtype) + + # 2. Conditional forward pass + velocity = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=cond_prompt_embeds, + timestep=timestep, + img_ids=img_ids_2d, + txt_ids=txt_ids_2d, + image_cond_hidden_states=image_cond_hidden_states, + image_cond_ids=image_cond_ids_2d, + target_timestep=target_timestep, + attention_kwargs=resolved_attention_kwargs, + return_dict=False, + )[0] + + # 3. Unconditional forward pass for true CFG + if do_classifier_free_guidance: + uncond_prompt_embeds = negative_prompt_embeds.to(dtype=transformer_dtype) + uncond_txt_ids_2d = ( + self._as_2d_position_ids(negative_text_ids, "negative_text_ids").to(latents.device) + if negative_text_ids is not None + else self._prepare_text_ids(uncond_prompt_embeds.shape[1], device=latents.device) + ) + neg_velocity = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=uncond_prompt_embeds, + timestep=timestep, + img_ids=img_ids_2d, + txt_ids=uncond_txt_ids_2d, + image_cond_hidden_states=image_cond_hidden_states, + image_cond_ids=image_cond_ids_2d, + target_timestep=target_timestep, + attention_kwargs=resolved_attention_kwargs, + return_dict=False, + )[0] + velocity = neg_velocity + guidance_scale * (velocity - neg_velocity) + + # 4. Scheduler step + output = self.scheduler.step( + velocity=velocity, + timestep=t, + latents=latents, + timestep_next=t_next, + next_latents=next_latents, + compute_log_prob=compute_log_prob, + return_dict=True, + return_kwargs=list(return_kwargs), + noise_level=noise_level, + ) + return output + + def forward( + self, + t: torch.Tensor, + latents: torch.Tensor, + latent_ids: Union[torch.Tensor, List[torch.Tensor]], + prompt_embeds: Union[torch.Tensor, List[torch.Tensor]], + text_ids: Union[torch.Tensor, List[torch.Tensor]], + image_latents: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + image_latent_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + negative_prompt_embeds: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + negative_text_ids: Optional[Union[torch.Tensor, List[Optional[torch.Tensor]]]] = None, + guidance_scale: float = 3.5, + t_next: Optional[torch.Tensor] = None, + next_latents: Optional[torch.Tensor] = None, + noise_level: Optional[float] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + compute_log_prob: bool = True, + return_kwargs: Sequence[str] = ( + "velocity", + "next_latents", + "next_latents_mean", + "std_dev_t", + "dt", + "log_prob", + ), + ) -> FlowMatchEulerDiscreteSDESchedulerOutput: + """General forward pass supporting T2I, multi-reference I2I, and ragged prompt windows.""" + has_image_cond = image_latents is not None and ( + isinstance(image_latents, torch.Tensor) + or any(item is not None for item in image_latents) + ) + has_ragged_prompts = ( + isinstance(prompt_embeds, list) + or isinstance(text_ids, list) + or isinstance(negative_prompt_embeds, list) + or isinstance(negative_text_ids, list) + ) + fall_back = has_image_cond or has_ragged_prompts + + if not fall_back: + return self._forward( + t=t, + latents=latents, + latent_ids=latent_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + image_latents=None, + image_latent_ids=None, + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + guidance_scale=guidance_scale, + t_next=t_next, + next_latents=next_latents, + noise_level=noise_level, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + compute_log_prob=compute_log_prob, + return_kwargs=return_kwargs, + ) + + if not self._has_warned_forward_fallback: + logger.warning( + "FLUX.3: Image conditioning or variable prompt window detected in forward(). " + "Processing samples individually (warning shown once)." + ) + self._has_warned_forward_fallback = True + + batch_size = latents.shape[0] + outputs: List[FlowMatchEulerDiscreteSDESchedulerOutput] = [] + + for idx in range(batch_size): + single_t = t[idx].unsqueeze(0) if t.ndim > 0 and t.shape[0] == batch_size else t + single_t_next = ( + t_next[idx].unsqueeze(0) + if t_next is not None and t_next.ndim > 0 and t_next.shape[0] == batch_size + else t_next + ) + single_latents = latents[idx].unsqueeze(0) + single_latent_ids = ( + latent_ids[idx].unsqueeze(0) + if isinstance(latent_ids, list) or latent_ids.ndim == 3 + else latent_ids + ) + single_next_latents = ( + next_latents[idx].unsqueeze(0) if next_latents is not None else None + ) + single_prompt_embeds = prompt_embeds[idx].unsqueeze(0) + single_text_ids = ( + text_ids[idx].unsqueeze(0) + if isinstance(text_ids, list) or text_ids.ndim == 3 + else text_ids + ) + single_image_latents = ( + image_latents[idx].unsqueeze(0) + if image_latents is not None and image_latents[idx] is not None + else None + ) + single_image_latent_ids = ( + image_latent_ids[idx].unsqueeze(0) + if image_latent_ids is not None and image_latent_ids[idx] is not None + else None + ) + single_neg_prompt_embeds = ( + negative_prompt_embeds[idx].unsqueeze(0) + if negative_prompt_embeds is not None and negative_prompt_embeds[idx] is not None + else None + ) + single_neg_text_ids = ( + negative_text_ids[idx].unsqueeze(0) + if negative_text_ids is not None and negative_text_ids[idx] is not None + else None + ) + + out = self._forward( + t=single_t, + latents=single_latents, + latent_ids=single_latent_ids, + prompt_embeds=single_prompt_embeds, + text_ids=single_text_ids, + image_latents=single_image_latents, + image_latent_ids=single_image_latent_ids, + negative_prompt_embeds=single_neg_prompt_embeds, + negative_text_ids=single_neg_text_ids, + guidance_scale=guidance_scale, + t_next=single_t_next, + next_latents=single_next_latents, + noise_level=noise_level, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + compute_log_prob=compute_log_prob, + return_kwargs=return_kwargs, + ) + outputs.append(out) + + outputs_dict = [o.to_dict() for o in outputs] + return FlowMatchEulerDiscreteSDESchedulerOutput.from_dict( + { + key: ( + torch.cat([item[key] for item in outputs_dict], dim=0) + if outputs_dict[0][key] is not None + else None + ) + for key in outputs_dict[0].keys() + } + ) diff --git a/src/flow_factory/models/registry.py b/src/flow_factory/models/registry.py index 6d42da5a..83e6a02d 100644 --- a/src/flow_factory/models/registry.py +++ b/src/flow_factory/models/registry.py @@ -30,6 +30,7 @@ "flux1-kontext": "flow_factory.models.flux.flux1_kontext.Flux1KontextAdapter", "flux2": "flow_factory.models.flux.flux2.Flux2Adapter", "flux2-klein": "flow_factory.models.flux.flux2_klein.Flux2KleinAdapter", + "flux3": "flow_factory.models.flux.flux3.Flux3Adapter", "qwen-image": "flow_factory.models.qwen_image.qwen_image.QwenImageAdapter", "qwen-image-edit-plus": "flow_factory.models.qwen_image.qwen_image_edit_plus.QwenImageEditPlusAdapter", "qwen-image-2.1": "flow_factory.models.qwen_image_21.qwen_image_21.QwenImage21Adapter", diff --git a/tests/docs/test_flux3_docs.py b/tests/docs/test_flux3_docs.py new file mode 100644 index 00000000..e71196b5 --- /dev/null +++ b/tests/docs/test_flux3_docs.py @@ -0,0 +1,29 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] + + +def test_readme_lists_flux3_klein_models() -> None: + readme = (ROOT / "README.md").read_text(encoding="utf-8") + for repo in ( + "https://huggingface.co/black-forest-labs/flux-3-image-klein-base", + "https://huggingface.co/black-forest-labs/flux-3-image-klein", + ): + row_start = readme.index(repo) + model_row = readme[row_start : readme.index("", row_start)] + assert "5.5B" in model_row + assert "flux3" in model_row diff --git a/tests/models/test_all_model_loaders.py b/tests/models/test_all_model_loaders.py index 31630831..9af909df 100644 --- a/tests/models/test_all_model_loaders.py +++ b/tests/models/test_all_model_loaders.py @@ -192,6 +192,36 @@ def test_bagel_adapter_imports_with_its_optional_kernel_contract( assert partial_inv_freq.shape == (2,) +def test_flux3_loader_builds_modular_pipeline_and_rewrites_local_specs( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + module = importlib.import_module("flow_factory.models.flux.flux3") + spec = SimpleNamespace( + pretrained_model_name_or_path="black-forest-labs/flux-3-image-klein-base", revision="main" + ) + pipeline = SimpleNamespace( + pretrained_component_names=["transformer"], + _component_specs={"transformer": spec}, + ) + calls = [] + + monkeypatch.setattr( + module.ModularPipeline, + "from_pretrained", + lambda path: calls.append(path) or pipeline, + ) + adapter = object.__new__(module.Flux3Adapter) + adapter.model_args = SimpleNamespace(model_name_or_path=str(tmp_path)) + + assert adapter.load_pipeline() is pipeline + assert calls == [str(tmp_path)] + assert spec.pretrained_model_name_or_path == str(tmp_path) + assert spec.revision is None + assert module.Flux3Adapter.supports_fsdp2_cpu_efficient_loading + assert module.Flux3Adapter.ddp_find_unused_parameters + + def test_all_registered_adapters_have_an_audited_load_pipeline() -> None: from flow_factory.models.registry import list_registered_models @@ -199,6 +229,7 @@ def test_all_registered_adapters_have_an_audited_load_pipeline() -> None: audited_paths = { f"{module_name}.{adapter_name}" for module_name, adapter_name in _CLASSIC_ADAPTERS } | { + "flow_factory.models.flux.flux3.Flux3Adapter", "flow_factory.models.bagel.bagel.BagelAdapter", "flow_factory.models.sensenova.sensenova.SenseNovaAdapter", "flow_factory.models.minimax_h3.adapters.MiniMaxH3T2VAAdapter", @@ -219,6 +250,7 @@ def test_model_type_literal_matches_registered_model_keys() -> None: def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: + from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.minimax_h3.adapters import ( MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter, @@ -239,6 +271,7 @@ def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: assert all( adapter.supports_fsdp2_cpu_efficient_loading for adapter in ( + Flux3Adapter, MiniMaxH3T2VAAdapter, MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter, @@ -249,6 +282,7 @@ def test_model_specific_load_dtype_defaults_are_explicit_and_narrow() -> None: "text_encoders": torch.bfloat16, "vae": torch.float32, } + assert dict(Flux3Adapter.component_load_dtype_defaults) == expected_wan_defaults assert Wan2_T2V_Adapter.component_load_dtype_defaults == expected_wan_defaults assert Wan2_I2V_Adapter.component_load_dtype_defaults == { **expected_wan_defaults, diff --git a/tests/models/test_flux3.py b/tests/models/test_flux3.py new file mode 100644 index 00000000..d01fb873 --- /dev/null +++ b/tests/models/test_flux3.py @@ -0,0 +1,712 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU contract, parity, and replay tests for the FLUX.3 adapter.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any, List + +import numpy as np +import pytest +import torch +from PIL import Image + +from diffusers import ( + AutoencoderKLFlux2, + ClassifierFreeGuidance, + FlowMatchEulerDiscreteScheduler, + Flux3Transformer2DModel, +) +from diffusers.modular_pipelines import PipelineState +from diffusers.modular_pipelines.flux3 import ( + Flux3AutoBlocks, + Flux3DistilledAutoBlocks, + Flux3DistilledModularPipeline, + Flux3ModularPipeline, +) +from diffusers.modular_pipelines.flux3.before_denoise import ( + Flux3DistilledTextRoPEInputsStep, + Flux3PrepareImageCondLatentsStep, + Flux3PrepareLatentsStep, + Flux3SetTimestepsStep, + Flux3TextRoPEInputsStep, +) +from diffusers.modular_pipelines.flux3.decoders import Flux3DecodeStep, Flux3UnpackLatentsStep +from diffusers.modular_pipelines.flux3.denoise import Flux3DenoiseStep, Flux3DistilledDenoiseStep +from diffusers.modular_pipelines.flux3.encoders import ( + Flux3ProcessImagesInputStep, + Flux3VaeEncoderStep, +) +from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor +from diffusers.modular_pipelines.flux3.inputs import ( + Flux3DistilledTextInputStep, + Flux3ImageCondInputStep, + Flux3TextInputStep, +) +from flow_factory.contracts import ( + BatchCapability, + GeometrySource, + InputMediaOrder, + NegativePromptPolicy, +) +from flow_factory.hparams import Arguments +from flow_factory.models.abc import BaseAdapter +from flow_factory.models.flux._output import encode_flux3_output_images +from flow_factory.models.flux.flux3 import ( + IMAGE_ROPE_SCALE, + TEXT_PAD_MULTIPLE, + Flux3Adapter, + Flux3Sample, +) +from flow_factory.models.registry import get_model_adapter_class +from flow_factory.models.runtime import ModularPipelineRuntime +from flow_factory.samples import BaseSample +from flow_factory.scheduler import FlowMatchEulerDiscreteSDEScheduler + + +class _FakeInnerTokenizer: + def __call__( + self, + text: str, + *, + return_tensors: str = "pt", + truncation: bool = True, + max_length: int = 512, + padding: str | bool = False, + padding_side: str = "right", + ) -> dict[str, torch.Tensor]: + assert return_tensors == "pt" + assert truncation is True + raw_tokens = [min(ord(ch), 250) for ch in text][:max_length] + if not raw_tokens: + raw_tokens = [1] + if padding == "max_length": + assert padding_side == "right" + pad_count = max(0, max_length - len(raw_tokens)) + input_ids = raw_tokens + [0] * pad_count + attention_mask = [1] * len(raw_tokens) + [0] * pad_count + else: + input_ids = raw_tokens + attention_mask = [1] * len(raw_tokens) + return { + "input_ids": torch.tensor([input_ids], dtype=torch.long), + "attention_mask": torch.tensor([attention_mask], dtype=torch.long), + } + + +class _FakeProcessor: + def __init__(self) -> None: + self.tokenizer = _FakeInnerTokenizer() + self.chat_calls: list[str] = [] + + def apply_chat_template( + self, + messages: list[dict[str, str]], + *, + tokenize: bool = False, + add_generation_prompt: bool = True, + ) -> str: + assert tokenize is False + assert add_generation_prompt is True + content = messages[0]["content"] + self.chat_calls.append(content) + return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" + + +class _FakeQwen3VL(torch.nn.Module): + def __init__(self, hidden_dim: int = 4, num_layers: int = 33) -> None: + super().__init__() + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.dummy_param = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + self.calls: list[dict[str, Any]] = [] + + @property + def dtype(self) -> torch.dtype: + return self.dummy_param.dtype + + @property + def device(self) -> torch.device: + return self.dummy_param.device + + def forward( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + output_hidden_states: bool = True, + use_cache: bool = False, + ) -> Any: + assert output_hidden_states is True + assert use_cache is False + self.calls.append( + { + "input_ids": input_ids.clone(), + "attention_mask": attention_mask.clone(), + "position_ids": position_ids.clone(), + } + ) + batch_size, seq_len = input_ids.shape + base = ( + input_ids.to(torch.float32) * attention_mask.to(torch.float32) + + position_ids.to(torch.float32) * 0.01 + ) + base = base.unsqueeze(-1).repeat(1, 1, self.hidden_dim) + hidden_states = tuple(base + float(layer_idx) * 0.1 for layer_idx in range(self.num_layers)) + return SimpleNamespace(hidden_states=hidden_states) + + +def _build_tiny_flux3_components(*, flow_map: bool = False) -> dict[str, Any]: + torch.manual_seed(0) + transformer = Flux3Transformer2DModel( + in_channels=16, + cond_in_channels=16, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=16, + timestep_channels=32, + axes_dims_rope=(2, 2, 2, 2), + flow_map=flow_map, + ).eval() + vae = AutoencoderKLFlux2( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"), + up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"), + block_out_channels=(8, 16), + layers_per_block=1, + latent_channels=4, + norm_num_groups=4, + ).eval() + image_processor = Flux3ImageProcessor(vae_scale_factor=4, vae_latent_channels=4) + scheduler = FlowMatchEulerDiscreteScheduler(shift=6.93, use_dynamic_shifting=False) + text_encoder = _FakeQwen3VL(hidden_dim=2, num_layers=33).eval() + tokenizer = _FakeProcessor() + + blocks_cls = Flux3DistilledAutoBlocks if flow_map else Flux3AutoBlocks + pipeline: Flux3ModularPipeline = blocks_cls().init_pipeline() + update_kwargs: dict[str, Any] = { + "transformer": transformer, + "vae": vae, + "image_processor": image_processor, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + if not flow_map: + update_kwargs["guider"] = ClassifierFreeGuidance(guidance_scale=3.5) + pipeline.update_components(**update_kwargs) + return { + "pipeline": pipeline, + "transformer": transformer, + "vae": vae, + "image_processor": image_processor, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + + +def _build_flux3_adapter(*, flow_map: bool = False, noise_level: float = 0.0) -> Flux3Adapter: + bundle = _build_tiny_flux3_components(flow_map=flow_map) + adapter = object.__new__(Flux3Adapter) + adapter.config = SimpleNamespace( + mixed_precision="no", + scheduler_args=SimpleNamespace( + noise_level=noise_level, + sde_steps=None, + num_sde_steps=None, + seed=42, + dynamics_type="Flow-SDE" if noise_level > 0 else "ODE", + ), + ) + adapter.model_args = SimpleNamespace( + model_name_or_path="tiny-flux3", + target_components=["transformer"], + target_modules="default", + finetune_type="lora", + trainable_parameters_dtype=torch.float32, + component_load_dtypes=None, + frozen_parameters_dtype=None, + ) + adapter.training_args = SimpleNamespace( + height=32, + width=32, + latent_storage_dtype=None, + enable_gradient_checkpointing=False, + ) + adapter.eval_args = SimpleNamespace() + adapter.accelerator = SimpleNamespace( + device=torch.device("cpu"), + is_main_process=True, + is_local_main_process=True, + process_index=0, + distributed_type=None, + mixed_precision="no", + unwrap_model=lambda model: model, + wait_for_everyone=lambda: None, + ) + adapter._mode = "eval" if noise_level == 0.0 else "rollout" + adapter._component_load_dtype_manifest = Flux3Adapter.component_load_dtype_defaults + adapter._component_load_dtype_overrides = None + adapter.component_runtime = ModularPipelineRuntime.from_adapter(adapter, bundle["pipeline"]) + adapter.pipeline = bundle["pipeline"] + adapter._components = adapter.component_runtime.override_components + adapter.pipeline.scheduler = FlowMatchEulerDiscreteSDEScheduler( + shift=6.93, + use_dynamic_shifting=False, + noise_level=noise_level, + dynamics_type="Flow-SDE" if noise_level > 0 else "ODE", + seed=42, + ) + adapter.scheduler_group = adapter.build_scheduler_group() + adapter._effective_pipeline_io_contract = adapter._resolve_pipeline_io_contract() + adapter._condition_state_preparer = adapter._build_condition_state_preparer_declaration() + adapter._condition_state_encoding_modules = ( + adapter._validate_condition_state_preparer_lifecycle() + ) + adapter._output_state_codec = adapter._build_output_state_codec_declaration() + adapter._output_state_encoding_modules = adapter._validate_output_state_codec_lifecycle() + adapter._has_warned_inference_fallback = False + adapter._has_warned_forward_fallback = False + return adapter + + +def test_registry_inheritance_and_io_contracts() -> None: + assert get_model_adapter_class("flux3") is Flux3Adapter + Flux3Adapter.validate_offline_output_capability() + assert BaseAdapter in Flux3Adapter.__bases__ + assert not any( + base.__module__.startswith("flow_factory.models.flux.") and base is not Flux3Adapter + for base in Flux3Adapter.__mro__ + ) + assert not Flux3Adapter.supports_diffusers_cache + assert Flux3Adapter.supports_fsdp2_cpu_efficient_loading + assert Flux3Adapter.ddp_find_unused_parameters + + base_adapter = _build_flux3_adapter(flow_map=False) + assert not base_adapter.is_distilled + base_contract = base_adapter.effective_pipeline_io_contract + assert base_contract is not None + assert base_contract.negative_prompt is NegativePromptPolicy.OPTIONAL + assert base_contract.geometry_source is GeometrySource.CONFIGURED + assert base_contract.batch_capability is BatchCapability.UNIFORM + assert base_contract.input_media.order is InputMediaOrder.WITHIN_TYPE + + distilled_adapter = _build_flux3_adapter(flow_map=True) + assert distilled_adapter.is_distilled + distilled_contract = distilled_adapter.effective_pipeline_io_contract + assert distilled_contract is not None + assert distilled_contract.negative_prompt is NegativePromptPolicy.UNSUPPORTED + with pytest.raises(ValueError, match="does not support negative_prompt"): + distilled_adapter.encode_prompt("a fox", negative_prompt="blurry", guidance_scale=1.0) + + +def test_grpo_example_parses_through_production_config() -> None: + root = Path(__file__).resolve().parents[2] + config = Arguments.load_from_yaml(str(root / "examples/grpo/lora/flux3/default.yaml")) + + assert config.model_args.model_type == "flux3" + assert config.model_args.model_name_or_path == "black-forest-labs/flux-3-image-klein-base" + assert config.training_args.trainer_type == "grpo" + assert config.training_args.guidance_scale == 3.5 + assert config.eval_args.guidance_scale == 3.5 + + +def test_prompt_encoding_enforces_80_token_windows_and_pad_position_ids() -> None: + adapter = _build_flux3_adapter(flow_map=False) + encoded = adapter.encode_prompt( + prompt="a red fox", + guidance_scale=3.5, + max_sequence_length=512, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + pad_multiple=TEXT_PAD_MULTIPLE, + ) + + assert isinstance(encoded["prompt_ids"], torch.Tensor) + assert isinstance(encoded["prompt_embeds"], torch.Tensor) + assert isinstance(encoded["text_ids"], torch.Tensor) + assert encoded["prompt_ids"].shape == (1, 80) + assert encoded["prompt_embeds"].shape == (1, 80, 16) + assert encoded["text_ids"].shape == (1, 80, 4) + assert encoded["negative_prompt_embeds"].shape == (1, 80, 16) + assert encoded["negative_text_ids"].shape == (1, 80, 4) + + text_encoder: _FakeQwen3VL = adapter.text_encoder # type: ignore[assignment] + first_call = text_encoder.calls[0] + pad_positions = first_call["position_ids"][first_call["attention_mask"] == 0] + assert pad_positions.numel() > 0 + assert torch.all(pad_positions == 1) + + # Variable-length prompts with different 80-token windows remain un-merged so shorter + # prompts are never polluted with extra unmasked padding tokens. + short_prompt = "fox" + long_prompt = "word " * 15 # 75 chars + chat template -> (80, 160] window + ragged = adapter.encode_prompt( + prompt=[short_prompt, long_prompt], + guidance_scale=1.0, + max_sequence_length=512, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + ) + assert isinstance(ragged["prompt_embeds"], list) + assert ragged["prompt_embeds"][0].shape == (80, 16) + assert ragged["prompt_embeds"][1].shape == (160, 16) + + +def test_condition_image_encoding_assigns_negative_l_coords_and_preserves_empty_slots() -> None: + adapter = _build_flux3_adapter(flow_map=False) + ref0 = Image.new("RGB", (16, 8), (120, 40, 200)) + ref1 = Image.new("RGB", (8, 8), (10, 180, 90)) + + encoded = adapter.encode_image([[], [ref0, ref1]], condition_image_size=(64, 64)) + + assert encoded["condition_images"][0] == [] + assert encoded["image_latents"][0] is None + assert encoded["image_latent_ids"][0] is None + + # Tiny VAE has scale_factor=2, so 2x2 patchify grid is 4px: + # ref0 (16x8) -> 4x2 = 8 tokens at l=-1; ref1 (8x8) -> 2x2 = 4 tokens at l=-2. + latents_1 = encoded["image_latents"][1] + ids_1 = encoded["image_latent_ids"][1] + assert isinstance(latents_1, torch.Tensor) + assert isinstance(ids_1, torch.Tensor) + assert latents_1.shape == (12, 16) + assert ids_1.shape == (12, 4) + assert torch.all(ids_1[:8, 3] == -1.0) + assert torch.all(ids_1[8:, 3] == -2.0) + assert torch.allclose(ids_1[:8, 1].max(), torch.tensor(1.0 * IMAGE_ROPE_SCALE)) + assert torch.allclose(ids_1[:8, 2].max(), torch.tensor(3.0 * IMAGE_ROPE_SCALE)) + + +def test_output_codec_and_decode_latents_round_trip() -> None: + adapter = _build_flux3_adapter(flow_map=False) + pixels = torch.zeros(2, 3, 16, 16) + encoded = encode_flux3_output_images( + adapter, + pixels, + torch.Generator().manual_seed(7), + image_rope_scale=IMAGE_ROPE_SCALE, + ) + assert encoded.latents.shape == (2, 16, 16) + assert encoded.forward_context["latent_ids"].shape == (2, 16, 4) + + decoded_from_hw = adapter.decode_latents( + encoded.latents, + height=16, + width=16, + output_type="pt", + ) + decoded_from_ids = adapter.decode_latents( + encoded.latents, + latent_ids=encoded.forward_context["latent_ids"], + output_type="pt", + ) + assert decoded_from_hw.shape == (2, 3, 16, 16) + torch.testing.assert_close(decoded_from_hw, decoded_from_ids, rtol=0, atol=0) + + +@pytest.mark.parametrize("flow_map", [False, True]) +def test_end_to_end_parity_against_diffusers_modular_blocks(flow_map: bool) -> None: + """Verify exact numerical parity between Flux3Adapter.inference() and Diffusers blocks.""" + adapter = _build_flux3_adapter(flow_map=flow_map, noise_level=0.0) + pipe = adapter.pipeline + height, width = 16, 16 + num_inference_steps = 3 + guidance_scale = 1.0 if flow_map else 3.5 + layers = (4, 8, 12, 16, 20, 24, 28, 32) + ref_images = [ + Image.fromarray(np.full((16, 8, 3), 80, dtype=np.uint8)), + Image.fromarray(np.full((8, 8, 3), 160, dtype=np.uint8)), + ] + + # 1. Encode prompts with adapter + prompt_enc = adapter.encode_prompt( + prompt="a snowy forest", + negative_prompt="" if not flow_map else None, + guidance_scale=guidance_scale, + max_sequence_length=64, + text_encoder_out_layers=layers, + ) + + # 2. Run official Diffusers modular blocks on the same components & seed + ref_pipe = (Flux3DistilledAutoBlocks if flow_map else Flux3AutoBlocks)().init_pipeline() + ref_update = { + "transformer": adapter.transformer, + "vae": adapter.vae, + "image_processor": pipe.image_processor, + "scheduler": FlowMatchEulerDiscreteScheduler(shift=6.93, use_dynamic_shifting=False), + "text_encoder": adapter.text_encoder, + "tokenizer": adapter.processor, + } + if not flow_map: + ref_update["guider"] = ClassifierFreeGuidance(guidance_scale=guidance_scale) + ref_pipe.update_components(**ref_update) + + state = PipelineState( + values={ + "prompt_embeds": prompt_enc["prompt_embeds"].clone(), + "negative_prompt_embeds": ( + None if flow_map else prompt_enc["negative_prompt_embeds"].clone() + ), + "image": ref_images, + "height": height, + "width": width, + "num_images_per_prompt": 1, + "generator": torch.Generator().manual_seed(123), + "num_inference_steps": num_inference_steps, + "timesteps": None, + "sigmas": None, + "latents": None, + "attention_kwargs": None, + "output_type": "pt", + } + ) + steps: List[Any] = [ + Flux3ProcessImagesInputStep(), + Flux3VaeEncoderStep(), + Flux3DistilledTextInputStep() if flow_map else Flux3TextInputStep(), + Flux3ImageCondInputStep(default_height=height, default_width=width), + Flux3PrepareLatentsStep( + default_height=height, + default_width=width, + image_rope_scale=IMAGE_ROPE_SCALE, + ), + Flux3PrepareImageCondLatentsStep(image_rope_scale=IMAGE_ROPE_SCALE), + Flux3SetTimestepsStep(shift=6.93, num_inference_steps=num_inference_steps), + Flux3DistilledTextRoPEInputsStep() if flow_map else Flux3TextRoPEInputsStep(), + Flux3DistilledDenoiseStep() if flow_map else Flux3DenoiseStep(), + Flux3UnpackLatentsStep(), + Flux3DecodeStep(), + ] + for step in steps: + ref_pipe, state = step(ref_pipe, state) + expected_images = state.values["images"] + + # 3. Run Flux3Adapter.inference() with the same seed and condition images + samples = adapter.inference( + images=[ref_images], + prompt=["a snowy forest"], + negative_prompt=[""] if not flow_map else None, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + condition_image_size=(64, 64), + max_sequence_length=64, + text_encoder_out_layers=layers, + generator=torch.Generator().manual_seed(123), + compute_log_prob=False, + ) + + assert len(samples) == 1 + torch.testing.assert_close(samples[0].image, expected_images[0], rtol=0, atol=1e-6) + + +def test_sde_rollout_replay_parity_and_gradient_flow() -> None: + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.7) + adapter.transformer.train() + ref_image = Image.fromarray(np.full((8, 8, 3), 100, dtype=np.uint8)) + + samples = adapter.inference( + images=[[ref_image]], + prompt=["relight "], + height=16, + width=16, + num_inference_steps=3, + guidance_scale=2.0, + condition_image_size=(64, 64), + max_sequence_length=64, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + generator=torch.Generator().manual_seed(99), + compute_log_prob=True, + ) + sample = samples[0] + assert isinstance(sample, Flux3Sample) + assert sample.all_latents is not None + assert sample.log_probs is not None + assert sample.all_latents.shape == (4, 16, 16) + assert sample.log_probs.shape == (2,) + + batch = BaseSample.stack(samples) + step_idx = 0 + t = batch["timesteps"][:, step_idx] + t_next = batch["timesteps"][:, step_idx + 1] + latents = batch["all_latents"][:, step_idx] + next_latents = batch["all_latents"][:, step_idx + 1] + noise_level = adapter.scheduler.get_noise_level_for_timestep(t[0]) + + replay = adapter.forward( + t=t, + t_next=t_next, + latents=latents, + next_latents=next_latents, + latent_ids=batch["latent_ids"], + prompt_embeds=batch["prompt_embeds"], + text_ids=batch["text_ids"], + image_latents=batch["image_latents"], + image_latent_ids=batch["image_latent_ids"], + negative_prompt_embeds=batch["negative_prompt_embeds"], + negative_text_ids=batch["negative_text_ids"], + guidance_scale=2.0, + noise_level=noise_level, + compute_log_prob=True, + ) + + assert replay.log_prob is not None + torch.testing.assert_close( + replay.log_prob.detach(), + batch["log_probs"][:, step_idx], + rtol=0, + atol=1e-6, + ) + ratio = torch.exp(replay.log_prob - batch["log_probs"][:, step_idx]) + torch.testing.assert_close(ratio.detach(), torch.ones_like(ratio), rtol=0, atol=1e-6) + + loss = -replay.log_prob.mean() + loss.backward() + + assert adapter.transformer.txt_mode_blocks[0].attn.to_q.weight.grad is not None + assert adapter.transformer.mode_blocks["image"][0].attn.to_q.weight.grad is not None + assert adapter.transformer.mode_blocks["image_cond"][0].attn.to_q.weight.grad is not None + assert adapter.transformer.single_transformer_blocks[0].attn.to_q.weight.grad is not None + + +def test_gradient_checkpointing_order_and_lora_round_trip(tmp_path: Path) -> None: + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.0) + units = adapter._gradient_checkpointing_units("transformer", adapter.transformer) + unit_names = [name for name, _ in units] + assert unit_names == [ + "txt_mode_blocks.0", + "mode_blocks.image.0", + "mode_blocks.image_cond.0", + "single_transformer_blocks.0", + ] + + adapter.model_args.lora_rank = 4 + adapter.model_args.lora_alpha = 8 + adapter.target_module_map = {"transformer": adapter.default_target_modules} + adapter.apply_lora( + target_modules=adapter.default_target_modules, + components=["transformer"], + ) + + trainable_params = adapter.get_trainable_parameters() + assert len(trainable_params) > 0 + with torch.no_grad(): + for param in trainable_params: + param.add_(0.05) + expected_snapshot = [param.detach().clone() for param in trainable_params] + + ckpt_dir = tmp_path / "flux3_lora_ckpt" + adapter.save_checkpoint(str(ckpt_dir), save_ema=False, model_only=True) + + with torch.no_grad(): + for param in trainable_params: + param.zero_() + adapter.load_checkpoint(str(ckpt_dir), resume_type="lora") + + restored_params = adapter.get_trainable_parameters() + for expected, restored in zip(expected_snapshot, restored_params, strict=True): + torch.testing.assert_close(restored, expected, rtol=0, atol=0) + + +def test_offline_sft_forward_state_and_load_plan_contracts() -> None: + from flow_factory.loading import ComponentRole, build_adapter_load_plan + + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.0) + adapter.training_args.height = 16 + adapter.training_args.width = 16 + + # 1. Verify ModelLoadCoordinator plan partitions TARGET, HOST, and AUXILIARY + plan = build_adapter_load_plan(adapter) + assert plan.request_for_component("transformer").role is ComponentRole.TARGET + assert plan.request_for_component("tokenizer").role is ComponentRole.HOST + assert plan.request_for_component("scheduler").role is ComponentRole.HOST + assert plan.request_for_component("text_encoder").role is ComponentRole.AUXILIARY + assert plan.request_for_component("vae").role is ComponentRole.AUXILIARY + assert plan.request_for_component("image_processor").role is ComponentRole.HOST + + # 2. Verify offline output-state codec + forward_state velocity matching + target_image = Image.new("RGB", (16, 16), (60, 120, 180)) + media_batch = ( + (SimpleNamespace(type="image", payload=target_image, fps=None, sample_rate=None),), + ) + prompt_enc = adapter.encode_prompt( + prompt="a castle", + guidance_scale=1.0, + max_sequence_length=64, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + ) + condition = { + "prompt_embeds": prompt_enc["prompt_embeds"], + "text_ids": prompt_enc["text_ids"], + } + prepared_condition = adapter.prepare_condition_state(condition) + encoded_output = adapter.encode_output_state( + media_batch, + prepared_condition, + torch.Generator().manual_seed(17), + ) + + times = adapter.build_training_component_times(torch.tensor([500.0])) + noised = adapter.add_forward_process_noise( + encoded_output.clean_state, + times, + generator=torch.Generator().manual_seed(23), + ) + forward_kwargs = { + **prepared_condition.condition, + **prepared_condition.forward_context, + **encoded_output.forward_context, + **dict(adapter.offline_training_forward_overrides), + } + step_out = adapter.forward_state( + batch=forward_kwargs, + state=noised.state, + times=times, + compute_log_prob=False, + return_fields=("velocity",), + ) + assert step_out.velocity is not None + per_elem_sq_error = { + name: (step_out.velocity.components[name] - noised.target_velocity.components[name]).pow(2) + for name in adapter.trajectory_component_order + } + reduced = adapter.reduce_flow_matching_objective_values(per_elem_sq_error, state=noised.state) + assert reduced.shape == (1,) + assert torch.isfinite(reduced).all() + + # 3. Verify Flux3Sample reconstruction_required_fields + sample = Flux3Sample( + prompt="a castle", + image=torch.zeros(3, 16, 16), + height=16, + width=16, + latent_ids=torch.zeros(16, 4), + text_ids=torch.zeros(80, 4), + ) + reconstructed = Flux3Sample( + **{ + field_name: getattr(sample, field_name) + for field_name in Flux3Sample.reconstruction_required_fields + if hasattr(sample, field_name) + } + ) + assert isinstance(reconstructed, Flux3Sample) diff --git a/tests/models/test_modern_image_output_codecs.py b/tests/models/test_modern_image_output_codecs.py index 16b8e6ff..d2711e2d 100644 --- a/tests/models/test_modern_image_output_codecs.py +++ b/tests/models/test_modern_image_output_codecs.py @@ -32,9 +32,11 @@ from flow_factory.models.flux._output import ( encode_flux2_vae_image, prepare_flux2_condition_latents, + prepare_flux3_condition_latents, ) from flow_factory.models.flux.flux2 import Flux2Adapter from flow_factory.models.flux.flux2_klein import Flux2KleinAdapter +from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.qwen_image._output import encode_qwen_vae_image from flow_factory.models.qwen_image.qwen_image import QwenImageAdapter from flow_factory.models.qwen_image.qwen_image_edit_plus import QwenImageEditPlusAdapter @@ -180,6 +182,7 @@ def _pack_qwen_latents( [ Flux2Adapter, Flux2KleinAdapter, + Flux3Adapter, QwenImageAdapter, QwenImageEditPlusAdapter, QwenImage21Adapter, @@ -452,6 +455,12 @@ def test_qwen_edit_keeps_condition_argmax_and_target_first_ordered_shapes() -> N BatchCapability.UNIFORM, InputMediaOrder.WITHIN_TYPE, ), + ( + Flux3Adapter, + GeometrySource.CONFIGURED, + BatchCapability.UNIFORM, + InputMediaOrder.WITHIN_TYPE, + ), ( QwenImageAdapter, GeometrySource.CONFIGURED, diff --git a/tests/models/test_offline_training_guidance.py b/tests/models/test_offline_training_guidance.py index c06d7637..f17ade16 100644 --- a/tests/models/test_offline_training_guidance.py +++ b/tests/models/test_offline_training_guidance.py @@ -19,6 +19,7 @@ from flow_factory.models.flux.flux1_kontext import Flux1KontextAdapter from flow_factory.models.flux.flux2 import Flux2Adapter from flow_factory.models.flux.flux2_klein import Flux2KleinAdapter +from flow_factory.models.flux.flux3 import Flux3Adapter from flow_factory.models.qwen_image.qwen_image import QwenImageAdapter from flow_factory.models.qwen_image.qwen_image_edit_plus import QwenImageEditPlusAdapter from flow_factory.models.qwen_image_21.qwen_image_21 import QwenImage21Adapter @@ -39,6 +40,7 @@ def test_offline_training_forward_overrides_are_immutable() -> None: BaseAdapter, SD3_5Adapter, Flux2KleinAdapter, + Flux3Adapter, QwenImageAdapter, QwenImageEditPlusAdapter, QwenImage21Adapter, From 28a98de25b1435597087af34b0509271f051cead Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 17 Oct 2026 16:22:42 +0800 Subject: [PATCH 2/8] [model] fix: default FLUX.3 encode_prompt to the official 8192-token limit Flow-Factory exposes no max_sequence_length setting, so dataset preprocessing encoded prompts with the encode_prompt default of 512 tokens, while the official FLUX.3 text encoder step defaults to 8192. A 999-token caption produced a 512-token window instead of the official 1040-token one: the prompt was truncated, and the window was no longer a multiple of the 80-token pad recipe. --- src/flow_factory/models/flux/flux3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py index 5e99f9ad..b462a6c5 100644 --- a/src/flow_factory/models/flux/flux3.py +++ b/src/flow_factory/models/flux/flux3.py @@ -427,7 +427,7 @@ def encode_prompt( negative_prompt: Optional[Union[str, List[str]]] = None, guidance_scale: float = 3.5, device: Optional[torch.device] = None, - max_sequence_length: int = 512, + max_sequence_length: int = 8192, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, pad_multiple: int = TEXT_PAD_MULTIPLE, ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]: From 8f790085167af6385f2fc2d3c6892da280d63784 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 17 Oct 2026 16:22:42 +0800 Subject: [PATCH 3/8] [dataset,test] fix: register flux3 in the offline smoke catalog and README table check Two existing tests failed once flux3 was registered: - test_offline_smoke_profiles requires every registered model type to appear in an offline smoke profile. Add flux3 to text_to_image, image_to_image, and multi_image_to_image, matching flux2-klein and the adapter's within-type multi-reference contract. - test_sensenova_docs locates the combined T2I/I2I README table by its rowspan, which grew from 9 to 11 with the two FLUX.3 rows. --- dataset/offline_smoke/profiles.py | 8 +++++--- tests/docs/test_sensenova_docs.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/dataset/offline_smoke/profiles.py b/dataset/offline_smoke/profiles.py index 5bb0c9db..14d0181c 100644 --- a/dataset/offline_smoke/profiles.py +++ b/dataset/offline_smoke/profiles.py @@ -230,9 +230,11 @@ def _case( _MODEL_TYPES = { - "text_to_image": "sd3-5 flux1 flux2 flux2-klein qwen-image qwen-image-2.1 z-image bagel sensenova", - "image_to_image": "flux1-kontext flux2 flux2-klein qwen-image-edit-plus bagel sensenova", - "multi_image_to_image": "flux2 flux2-klein qwen-image-edit-plus bagel sensenova", + "text_to_image": ( + "sd3-5 flux1 flux2 flux2-klein flux3 qwen-image qwen-image-2.1 z-image bagel sensenova" + ), + "image_to_image": "flux1-kontext flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova", + "multi_image_to_image": "flux2 flux2-klein flux3 qwen-image-edit-plus bagel sensenova", "text_to_video": "wan2_t2v", "first_frame_to_video": "wan2_i2v", "first_last_frame_to_video": "wan2_i2v", diff --git a/tests/docs/test_sensenova_docs.py b/tests/docs/test_sensenova_docs.py index abedc4a1..1c3a0e38 100644 --- a/tests/docs/test_sensenova_docs.py +++ b/tests/docs/test_sensenova_docs.py @@ -24,7 +24,7 @@ def _text(path: str) -> str: def test_main_readme_lists_sensenova_as_t2i_and_multi_reference_i2i() -> None: readme = _text("README.md") t2i_start = readme.index('Text-to-Image') - combined_start = readme.index('Text-to-Image & Image(s)-to-Image') + combined_start = readme.index('Text-to-Image & Image(s)-to-Image') video_start = readme.index('Text-to-Video') t2i_only_rows = readme[t2i_start:combined_start] From 26c7956046ffe58bed54317ae2f7c9fa0eb7523c Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 17 Oct 2026 16:24:22 +0800 Subject: [PATCH 4/8] [model,docs,agents] fix: align Diffusers submodule references with the 7ecb1e1b pin The FLUX.3 change moves the shared Diffusers submodule from 80c7ed262 to 7ecb1e1b. That pin raises the submodule's floor to Python 3.11 and huggingface-hub 1.32, and Qwen-Image 2.1 still named the old commit in its install error. - Point QWEN_IMAGE_21_DIFFUSERS_COMMIT at 7ecb1e1b. The functions it documents as copied (calculate_dimensions, retrieve_latents, _get_qwen_prompt_embeds, _pack_latents, _unpack_latents) are unchanged between the two commits, and the released checkpoint sets no sample_sigmas, so Qwen-Image 2.1 sampling is unaffected. - State the Python 3.11 requirement in the README install note and in dependencies.md. --- .agents/knowledge/dependencies.md | 2 ++ README.md | 3 ++- src/flow_factory/models/qwen_image_21/output_codec.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md index c84df288..4b5d39ac 100644 --- a/.agents/knowledge/dependencies.md +++ b/.agents/knowledge/dependencies.md @@ -80,6 +80,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package ### diffusers - Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule may be used for upstream development, but must not silently override the declared runtime dependency. +- `flux3` and `qwen-image-2.1` need the pinned submodule (`7ecb1e1b`), which requires Python 3.11+ + and `huggingface-hub>=1.32`; released diffusers keeps Python 3.10 support for the other adapters. - Model adapters depend on specific pipeline classes from diffusers. Major version bumps may rename or remove pipeline classes. - `load_pipeline()` in each adapter returns a `DiffusionPipeline`-compatible object; breaking changes in diffusers' pipeline API require adapter updates. diff --git a/README.md b/README.md index 17501764..f23bb6e5 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,8 @@ * **[2026-10-16]** **FLUX.3 Image Klein** support! Train both the base (`black-forest-labs/flux-3-image-klein-base`) and step-distilled flow-map (`black-forest-labs/flux-3-image-klein`) text-to-image and multi-reference image-to-image models with the [GRPO + LoRA recipe](examples/grpo/lora/flux3/default.yaml). - Until FLUX.3 reaches a released Diffusers package, install the pinned submodule: + Until FLUX.3 reaches a released Diffusers package, install the pinned submodule + (it requires Python 3.11 or newer): ```bash git submodule update --init pip install -e ./diffusers diff --git a/src/flow_factory/models/qwen_image_21/output_codec.py b/src/flow_factory/models/qwen_image_21/output_codec.py index 8113ae22..3a5fd2b1 100644 --- a/src/flow_factory/models/qwen_image_21/output_codec.py +++ b/src/flow_factory/models/qwen_image_21/output_codec.py @@ -23,7 +23,7 @@ from ..configured_image_output import EncodedImageTensor -QWEN_IMAGE_21_DIFFUSERS_COMMIT = "80c7ed262aeffbeb43ef13ae04baeb9b84515a69" +QWEN_IMAGE_21_DIFFUSERS_COMMIT = "7ecb1e1b4fe500be7fb7381a7bb79a24d7fa3941" def _retrieve_vae_latents( From 3f4cd051c913ea6dd102551a2461c9cfc8476e41 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 17 Oct 2026 21:04:45 +0800 Subject: [PATCH 5/8] [model,test] fix: default every FLUX.3 entry point to the official 8192 tokens encode_prompt already defaults to the official 8192-token limit, but inference, _inference, and the Qwen3-VL helpers kept their own 512 default and pass it to encode_prompt explicitly, so raw-prompt rollouts (for example with enable_preprocess: false) truncated long captions to 512 tokens. Share one DEFAULT_MAX_SEQUENCE_LENGTH across every entry point, run the window tests on the default, and cover the inference path. --- src/flow_factory/models/flux/flux3.py | 12 +++++---- tests/models/test_flux3.py | 39 ++++++++++++++++++++++++--- 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py index b462a6c5..42b77539 100644 --- a/src/flow_factory/models/flux/flux3.py +++ b/src/flow_factory/models/flux/flux3.py @@ -97,6 +97,8 @@ CONDITION_IMAGE_SIZE = (1024, 1024) IMAGE_ROPE_SCALE = 0.5 TEXT_PAD_MULTIPLE = 80 +# Matches the `Flux3TextEncoderStep` inference default, so long captions are not truncated. +DEFAULT_MAX_SEQUENCE_LENGTH = 8192 DEFAULT_TEXT_ENCODER_OUT_LAYERS: Tuple[int, ...] = (4, 8, 12, 16, 20, 24, 28, 32) @@ -300,7 +302,7 @@ def _encode_single_qwen3_vl_prompt( *, dtype: torch.dtype, device: torch.device, - max_sequence_length: int = 512, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, pad_multiple: int = TEXT_PAD_MULTIPLE, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -377,7 +379,7 @@ def _get_qwen3_vl_prompt_embeds( *, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, - max_sequence_length: int = 512, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, pad_multiple: int = TEXT_PAD_MULTIPLE, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, ) -> Tuple[ @@ -427,7 +429,7 @@ def encode_prompt( negative_prompt: Optional[Union[str, List[str]]] = None, guidance_scale: float = 3.5, device: Optional[torch.device] = None, - max_sequence_length: int = 8192, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, pad_multiple: int = TEXT_PAD_MULTIPLE, ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]: @@ -827,7 +829,7 @@ def _inference( condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, attention_kwargs: Optional[Dict[str, Any]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, - max_sequence_length: int = 512, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, pad_multiple: int = TEXT_PAD_MULTIPLE, compute_log_prob: bool = False, @@ -1087,7 +1089,7 @@ def inference( condition_image_size: Union[int, Tuple[int, int]] = CONDITION_IMAGE_SIZE, attention_kwargs: Optional[Dict[str, Any]] = None, joint_attention_kwargs: Optional[Dict[str, Any]] = None, - max_sequence_length: int = 512, + max_sequence_length: int = DEFAULT_MAX_SEQUENCE_LENGTH, text_encoder_out_layers: Sequence[int] = DEFAULT_TEXT_ENCODER_OUT_LAYERS, pad_multiple: int = TEXT_PAD_MULTIPLE, prompt_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, diff --git a/tests/models/test_flux3.py b/tests/models/test_flux3.py index d01fb873..913b7d9b 100644 --- a/tests/models/test_flux3.py +++ b/tests/models/test_flux3.py @@ -16,6 +16,7 @@ from __future__ import annotations +import math from pathlib import Path from types import SimpleNamespace from typing import Any, List @@ -49,6 +50,7 @@ from diffusers.modular_pipelines.flux3.denoise import Flux3DenoiseStep, Flux3DistilledDenoiseStep from diffusers.modular_pipelines.flux3.encoders import ( Flux3ProcessImagesInputStep, + Flux3TextEncoderStep, Flux3VaeEncoderStep, ) from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor @@ -67,6 +69,7 @@ from flow_factory.models.abc import BaseAdapter from flow_factory.models.flux._output import encode_flux3_output_images from flow_factory.models.flux.flux3 import ( + DEFAULT_MAX_SEQUENCE_LENGTH, IMAGE_ROPE_SCALE, TEXT_PAD_MULTIPLE, Flux3Adapter, @@ -85,7 +88,7 @@ def __call__( *, return_tensors: str = "pt", truncation: bool = True, - max_length: int = 512, + max_length: int | None = None, padding: str | bool = False, padding_side: str = "right", ) -> dict[str, torch.Tensor]: @@ -335,7 +338,6 @@ def test_prompt_encoding_enforces_80_token_windows_and_pad_position_ids() -> Non encoded = adapter.encode_prompt( prompt="a red fox", guidance_scale=3.5, - max_sequence_length=512, text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), pad_multiple=TEXT_PAD_MULTIPLE, ) @@ -362,7 +364,6 @@ def test_prompt_encoding_enforces_80_token_windows_and_pad_position_ids() -> Non ragged = adapter.encode_prompt( prompt=[short_prompt, long_prompt], guidance_scale=1.0, - max_sequence_length=512, text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), ) assert isinstance(ragged["prompt_embeds"], list) @@ -370,6 +371,38 @@ def test_prompt_encoding_enforces_80_token_windows_and_pad_position_ids() -> Non assert ragged["prompt_embeds"][1].shape == (160, 16) +def test_default_max_sequence_length_matches_upstream_and_keeps_long_prompts() -> None: + upstream_default = next( + param.default + for param in Flux3TextEncoderStep().inputs + if param.name == "max_sequence_length" + ) + assert DEFAULT_MAX_SEQUENCE_LENGTH == upstream_default == 8192 + + adapter = _build_flux3_adapter(flow_map=False) + # The fake tokenizer emits one token per character. + encoded = adapter.encode_prompt(prompt="x" * 600, guidance_scale=1.0) + + text_encoder: _FakeQwen3VL = adapter.text_encoder # type: ignore[assignment] + real_length = int(text_encoder.calls[-1]["attention_mask"].sum()) + assert real_length > 600 + expected_window = math.ceil(real_length / TEXT_PAD_MULTIPLE) * TEXT_PAD_MULTIPLE + assert encoded["prompt_embeds"].shape == (1, expected_window, 16) + assert encoded["text_ids"].shape == (1, expected_window, 4) + + # Raw-prompt rollouts (e.g. enable_preprocess: false) encode through inference(). + adapter.inference( + prompt=["x" * 600], + height=16, + width=16, + num_inference_steps=1, + guidance_scale=1.0, + generator=torch.Generator().manual_seed(0), + compute_log_prob=False, + ) + assert int(text_encoder.calls[-1]["attention_mask"].sum()) == real_length + + def test_condition_image_encoding_assigns_negative_l_coords_and_preserves_empty_slots() -> None: adapter = _build_flux3_adapter(flow_map=False) ref0 = Image.new("RGB", (16, 8), (120, 40, 200)) From ebac917ab45aee5eb792fe6c671f8737f34afd49 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sat, 17 Oct 2026 22:10:40 +0800 Subject: [PATCH 6/8] [model] docs: describe the FLUX.3 tokenizer without reward decoding #249 removes the reward-side prompt_ids decoding, so the tokenizer property no longer serves it. Match the wording #249 uses for the Qwen-Image 2.1 tokenizer. --- src/flow_factory/models/flux/flux3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py index 42b77539..24a5bfcb 100644 --- a/src/flow_factory/models/flux/flux3.py +++ b/src/flow_factory/models/flux/flux3.py @@ -208,7 +208,7 @@ def inference_modules(self) -> List[str]: @property def tokenizer(self) -> Any: - """Expose the processor-owned tokenizer to distributed reward decoding.""" + """Expose the tokenizer owned by the Qwen3-VL processor.""" if self.get_component("tokenizer") is None and self.has_component("tokenizer"): self.component_runtime.materialize_components(["tokenizer"]) processor = self.get_component("tokenizer") From d7994e1f008c429ceeb856b38cc7761e280e116a Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sun, 18 Oct 2026 07:52:23 +0800 Subject: [PATCH 7/8] [model,test,agents] fix: run every FLUX.3 sample at transformer batch size 1 for exact replay With per_device_batch_size > 1 the on-policy log-probs did not match the rollout: an exact probe found up to 7e-6 differences in 421 of 768 replays with FLUX.3-upsampled GenEval prompts and 140 of 768 with PickScore captions. inference() and forward() each ran a batch batched when all its prompts shared one 80-token window count and per sample otherwise. Group-contiguous rollout packs repeat one prompt and ran at batch size 8, while the shuffled optimize micro-batches mix window counts and replayed the same samples at batch size 1, and bf16 GEMMs of different shapes round differently. Both phases now run each sample's transformer forward at batch size 1, so a sample's result no longer depends on its batch, and _inference handles exactly one sample. A regression test regroups rollout packs into new micro-batches and requires bit-identical log-probs from batch size 1 transformer calls. The train-inference consistency notes record the failure mode. --- .../topics/train_inference_consistency.md | 9 ++ src/flow_factory/models/flux/flux3.py | 151 ++++-------------- tests/models/test_flux3.py | 56 ++++++- 3 files changed, 96 insertions(+), 120 deletions(-) diff --git a/.agents/knowledge/topics/train_inference_consistency.md b/.agents/knowledge/topics/train_inference_consistency.md index 3bbc5eaa..85c8882e 100644 --- a/.agents/knowledge/topics/train_inference_consistency.md +++ b/.agents/knowledge/topics/train_inference_consistency.md @@ -48,6 +48,7 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch boundary so mask/layout binding cannot drift between candidates. 11. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer. 12. **Different forward graph for the same inputs**: A prefix KV cache (prefill + cached decode) and a joint forward are mathematically equal but not bf16-equal. On Qwen-Image 2.1 they differ by 0.9–2.5% relative L2 in velocity (about 6e-6 in fp32), far outside `clip_range` and exact replay tolerances; grad mode alone does not change the result. Rollout and replay must build the same graph, so Qwen-Image 2.1 records `use_kv_cache` on the rollout sample and replay reads it back as a batch key. +13. **Execution path chosen from the batch's contents**: An adapter that runs a uniform batch batched and a mixed batch one sample at a time computes the same sample with different GEMM shapes whenever its rollout pack and its training micro-batch differ in that respect, and bf16 rounding then differs. FLUX.3 did this with its 80-token prompt windows: group-contiguous rollout packs repeat one prompt and ran batched, while shuffled micro-batches mix window counts and ran per sample, so about half the replays differed by up to 7e-6 in log-prob. Prompts long enough to span two windows trigger it, including some short captions once the chat template is added. FLUX.3 now runs every sample at transformer batch size 1 in both phases. ## Pack-composition-dependent adapters: `shuffle_samples` @@ -84,6 +85,14 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch - **Lesson**: Never flip the grad mode above code that decides from it whether a graph may outlive the call. A graph built under forced grad without a later backward is a leak whenever activation checkpointing saves non-tensor inputs that point back into it. - **Related Constraint**: N/A +### Run every FLUX.3 sample at transformer batch size 1 +- **Date**: 2026-10-18 +- **Symptom**: FLUX.3 GRPO with `per_device_batch_size: 8` logged on-policy ratios that rounded to 1.0000 but were not exactly 1 (`ratio_std` up to 1e-4), while Qwen-Image 2.1 under the same trainer stayed exactly 1. An exact-replay probe found 421 of 768 replayed log-probs off by up to 7e-6 with GenEval prompts rewritten by the FLUX.3 upsampler, and 140 of 768 with PickScore captions. +- **Root Cause**: `inference()` and `forward()` each chose batched execution for a batch whose prompts shared one 80-token window count and per-sample execution otherwise. Rollout packs of one repeated prompt ran at batch size 8; after the optimize-loop shuffle the same samples replayed at batch size 1, or the other way round. Keeping the rollout packs (`shuffle_samples: false`) made every replay exact, which isolated the pack composition. +- **Fix**: Both phases run each sample's transformer forward at batch size 1, so a sample's computation no longer depends on its batch. The regression test regroups rollout packs into new micro-batches and requires bit-identical log-probs and batch size 1 for every transformer call. +- **Lesson**: A GPU parity gate must include `per_device_batch_size > 1`, prompts of mixed lengths, and the default shuffle, and compare log-probs exactly rather than through the rounded ratio metric. Gates inherited from a `per_device_batch_size: 1` example never ran the batched path. +- **Related Constraint**: Constraint #7 + ## Cross-refs - UP: [`constraints.md` #7](../constraints.md#7-coupled-vs-decoupled-paradigm), [Architecture Execution Pipelines](../architecture.md#execution-pipelines) diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py index 24a5bfcb..474e1f4f 100644 --- a/src/flow_factory/models/flux/flux3.py +++ b/src/flow_factory/models/flux/flux3.py @@ -115,7 +115,12 @@ class Flux3Sample(I2ISample): class Flux3Adapter(ConfiguredImageOutputAdapterMixin, BaseAdapter): - """Adapter for FLUX.3 modular text-to-image and multi-reference image-to-image models.""" + """Adapter for FLUX.3 modular text-to-image and multi-reference image-to-image models. + + Rollout and replay run the transformer on one sample at a time. A sample's bf16 output then + does not depend on which samples share its rollout pack or training micro-batch, so the + on-policy ratio is exactly 1 for any ``per_device_batch_size`` and optimize-loop shuffle. + """ supports_diffusers_cache: ClassVar[bool] = False supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = True @@ -145,8 +150,6 @@ def __init__(self, config: Arguments, accelerator: Accelerator): super().__init__(config, accelerator) self.pipeline: Flux3ModularPipeline self.scheduler: FlowMatchEulerDiscreteSDEScheduler - self._has_warned_inference_fallback = False - self._has_warned_forward_fallback = False def load_pipeline(self) -> Flux3ModularPipeline: """Load the FLUX.3 modular pipeline specification from a local path or Hub repo.""" @@ -836,7 +839,7 @@ def _inference( extra_call_back_kwargs: List[str] = [], trajectory_indices: TrajectoryIndicesType = "all", ) -> List[Flux3Sample]: - """Run FLUX.3 rollout on a uniform-window T2I batch or a single sample.""" + """Run FLUX.3 rollout for a single sample.""" device = self.device transformer = self._require_component("transformer") dtype = transformer.dtype @@ -890,14 +893,19 @@ def _inference( or not isinstance(text_ids, torch.Tensor) ): raise ValueError( - "Flux3Adapter._inference requires uniform stacked prompt tensors; " - "use Flux3Adapter.inference() for variable-window batches." + "Flux3Adapter._inference requires the stacked prompt tensors of one sample; " + "use Flux3Adapter.inference() for batches." ) prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) if isinstance(negative_prompt_embeds, torch.Tensor): negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) batch_size = prompt_embeds.shape[0] + if batch_size != 1: + raise ValueError( + f"Flux3Adapter._inference runs one sample, received batch size {batch_size}; " + "use Flux3Adapter.inference() for batches." + ) # 2. Encode condition images images = ( @@ -1105,82 +1113,25 @@ def inference( extra_call_back_kwargs: List[str] = [], trajectory_indices: TrajectoryIndicesType = "all", ) -> List[Flux3Sample]: - """Run FLUX.3 inference; falls back to per-sample execution for I2I or ragged windows.""" + """Run FLUX.3 inference one sample at a time. + + Each sample's transformer forward has batch size 1, the same graph ``forward()`` replays + regardless of how training regroups samples into micro-batches. A batched forward would + round differently in bf16 whenever the rollout pack and the training micro-batch differ, + which they do under the optimize-loop shuffle or when prompts span different numbers of + 80-token windows. + """ if isinstance(prompt, str): prompt = [prompt] if isinstance(negative_prompt, str): negative_prompt = [negative_prompt] * (len(prompt) if prompt is not None else 1) - - is_nested_images = self._is_multi_images_batch(images) - is_nested_image_latents = self._is_multi_image_latents(image_latents) - is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( - negative_prompt_embeds, list - ) - - # If raw prompts are passed with batch_size > 1 and no pre-encoded embeds, check whether - # their 80-token windows differ by encoding once upfront. - if ( - prompt is not None - and len(prompt) > 1 - and prompt_embeds is None - and not (is_nested_images or is_nested_image_latents) - ): - encoded = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - guidance_scale=guidance_scale, - device=self.device, - max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, - pad_multiple=pad_multiple, - ) - prompt_ids = encoded["prompt_ids"] - prompt_embeds = encoded["prompt_embeds"] - text_ids = encoded["text_ids"] - negative_prompt_ids = encoded.get("negative_prompt_ids") - negative_prompt_embeds = encoded.get("negative_prompt_embeds") - negative_text_ids = encoded.get("negative_text_ids") - is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( - negative_prompt_embeds, list - ) - - fall_back = is_nested_images or is_nested_image_latents or is_ragged_prompts - if not fall_back: - return self._inference( - images=images, - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - condition_image_size=condition_image_size, - attention_kwargs=attention_kwargs, - joint_attention_kwargs=joint_attention_kwargs, - max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, - pad_multiple=pad_multiple, - prompt_ids=prompt_ids, # type: ignore[arg-type] - prompt_embeds=prompt_embeds, # type: ignore[arg-type] - text_ids=text_ids, # type: ignore[arg-type] - negative_prompt_ids=negative_prompt_ids, # type: ignore[arg-type] - negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] - negative_text_ids=negative_text_ids, # type: ignore[arg-type] - condition_images=condition_images, - image_latents=image_latents, # type: ignore[arg-type] - image_latent_ids=image_latent_ids, # type: ignore[arg-type] - compute_log_prob=compute_log_prob, - extra_call_back_kwargs=extra_call_back_kwargs, - trajectory_indices=trajectory_indices, - ) - - if not self._has_warned_inference_fallback: - logger.warning( - "FLUX.3 per-sample inference fallback activated (image conditioning or " - "variable 80-token prompt windows). This warning will only appear once." - ) - self._has_warned_inference_fallback = True + # A flat image list or unbatched latents hold the references of a single sample. + if images is not None and not self._is_multi_images_batch(images): + images = [images] + if isinstance(image_latents, torch.Tensor) and image_latents.ndim == 2: + image_latents = image_latents.unsqueeze(0) + if isinstance(image_latent_ids, torch.Tensor) and image_latent_ids.ndim == 2: + image_latent_ids = image_latent_ids.unsqueeze(0) if images is not None: batch_size = len(images) @@ -1444,47 +1395,11 @@ def forward( "log_prob", ), ) -> FlowMatchEulerDiscreteSDESchedulerOutput: - """General forward pass supporting T2I, multi-reference I2I, and ragged prompt windows.""" - has_image_cond = image_latents is not None and ( - isinstance(image_latents, torch.Tensor) - or any(item is not None for item in image_latents) - ) - has_ragged_prompts = ( - isinstance(prompt_embeds, list) - or isinstance(text_ids, list) - or isinstance(negative_prompt_embeds, list) - or isinstance(negative_text_ids, list) - ) - fall_back = has_image_cond or has_ragged_prompts - - if not fall_back: - return self._forward( - t=t, - latents=latents, - latent_ids=latent_ids, # type: ignore[arg-type] - prompt_embeds=prompt_embeds, # type: ignore[arg-type] - text_ids=text_ids, # type: ignore[arg-type] - image_latents=None, - image_latent_ids=None, - negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] - negative_text_ids=negative_text_ids, # type: ignore[arg-type] - guidance_scale=guidance_scale, - t_next=t_next, - next_latents=next_latents, - noise_level=noise_level, - attention_kwargs=attention_kwargs, - joint_attention_kwargs=joint_attention_kwargs, - compute_log_prob=compute_log_prob, - return_kwargs=return_kwargs, - ) - - if not self._has_warned_forward_fallback: - logger.warning( - "FLUX.3: Image conditioning or variable prompt window detected in forward(). " - "Processing samples individually (warning shown once)." - ) - self._has_warned_forward_fallback = True + """Replay one step for T2I and multi-reference I2I, one sample at a time. + Every sample runs through ``_forward`` with batch size 1, as in ``inference()``, so the + replayed log-probs match the rollout exactly whatever the micro-batch composition. + """ batch_size = latents.shape[0] outputs: List[FlowMatchEulerDiscreteSDESchedulerOutput] = [] diff --git a/tests/models/test_flux3.py b/tests/models/test_flux3.py index 913b7d9b..508d7111 100644 --- a/tests/models/test_flux3.py +++ b/tests/models/test_flux3.py @@ -287,8 +287,6 @@ def _build_flux3_adapter(*, flow_map: bool = False, noise_level: float = 0.0) -> ) adapter._output_state_codec = adapter._build_output_state_codec_declaration() adapter._output_state_encoding_modules = adapter._validate_output_state_codec_lifecycle() - adapter._has_warned_inference_fallback = False - adapter._has_warned_forward_fallback = False return adapter @@ -621,6 +619,60 @@ def test_sde_rollout_replay_parity_and_gradient_flow() -> None: assert adapter.transformer.single_transformer_blocks[0].attn.to_q.weight.grad is not None +def test_on_policy_replay_is_exact_for_any_micro_batch_composition() -> None: + adapter = _build_flux3_adapter(flow_map=False, noise_level=0.7) + adapter.transformer.train() + transformer_batch_sizes: list[int] = [] + adapter.transformer.register_forward_pre_hook( + lambda module, args, kwargs: transformer_batch_sizes.append( + kwargs["hidden_states"].shape[0] + ), + with_kwargs=True, + ) + short_prompt, long_prompt = "fox", "word " * 15 # one and two 80-token windows + rollout_kwargs = dict( + height=16, + width=16, + num_inference_steps=3, + guidance_scale=2.0, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + compute_log_prob=True, + ) + # A uniform pack, as group-contiguous sampling produces, and a pack of mixed window counts. + samples = adapter.inference( + prompt=[short_prompt, short_prompt], + generator=torch.Generator().manual_seed(7), + **rollout_kwargs, + ) + adapter.inference( + prompt=[short_prompt, long_prompt], + generator=torch.Generator().manual_seed(8), + **rollout_kwargs, + ) + assert set(transformer_batch_sizes) == {1} + + # Optimize micro-batches regroup the rollout packs, as the optimize-loop shuffle does. + for order in ([0, 2], [1, 3], [3, 2, 1, 0]): + batch = BaseSample.stack([samples[index] for index in order]) + for step_idx in range(2): + t = batch["timesteps"][:, step_idx] + replay = adapter.forward( + t=t, + t_next=batch["timesteps"][:, step_idx + 1], + latents=batch["all_latents"][:, step_idx], + next_latents=batch["all_latents"][:, step_idx + 1], + latent_ids=batch["latent_ids"], + prompt_embeds=batch["prompt_embeds"], + text_ids=batch["text_ids"], + negative_prompt_embeds=batch["negative_prompt_embeds"], + negative_text_ids=batch["negative_text_ids"], + guidance_scale=2.0, + noise_level=adapter.scheduler.get_noise_level_for_timestep(t[0]), + compute_log_prob=True, + ) + assert torch.equal(replay.log_prob.detach(), batch["log_probs"][:, step_idx]) + assert set(transformer_batch_sizes) == {1} + + def test_gradient_checkpointing_order_and_lora_round_trip(tmp_path: Path) -> None: adapter = _build_flux3_adapter(flow_map=False, noise_level=0.0) units = adapter._gradient_checkpointing_units("transformer", adapter.transformer) From cf12271ca967349680165044dccdea58617e7e09 Mon Sep 17 00:00:00 2001 From: Jayce-Ping <315229706@qq.com> Date: Sun, 18 Oct 2026 08:11:24 +0800 Subject: [PATCH 8/8] [trainer,model,test,docs,agents] feat: shuffle whole rollout packs for pack-dependent adapters train.shuffle_samples permuted individual samples, so an adapter whose batched forward depends on batch composition replayed samples in micro-batches that differed from their rollout packs, and the on-policy ratio was no longer exactly 1. Bagel could only warn and ask for shuffle_samples: false, and FLUX.3 ran every sample alone, which cost about 13% rollout and 15% optimization time at per_device_batch_size 8. For adapters that declare requires_preserved_replay_batch_composition, the generation loop now records the rollout packs and the optimize loop permutes whole packs, keeping each pack's samples and order. It fails fast when a pack does not match per_device_batch_size or the samples changed after rollout; reward overlap shares that validation. Batch-invariant adapters still shuffle individual samples with the same permutation, and shuffle_samples: false still keeps the rollout order. FLUX.3 declares the flag and goes back to batched execution for packs whose prompts share one 80-token window count. Bagel's warning is dropped, and the help text, guidance, and train-inference consistency notes describe the semantics. --- .../topics/train_inference_consistency.md | 19 +- guidance/algorithms.md | 2 +- guidance/gpu_validation.md | 2 +- guidance/workflow.md | 2 +- .../hparams/training_args/_base.py | 12 +- src/flow_factory/models/abc.py | 8 +- src/flow_factory/models/bagel/bagel.py | 10 - src/flow_factory/models/flux/flux3.py | 153 ++++++++++++--- src/flow_factory/trainers/abc.py | 99 +++++++--- tests/models/test_flux3.py | 46 +++-- tests/trainers/test_optimize_sample_order.py | 181 ++++++++++++++++++ 11 files changed, 423 insertions(+), 111 deletions(-) create mode 100644 tests/trainers/test_optimize_sample_order.py diff --git a/.agents/knowledge/topics/train_inference_consistency.md b/.agents/knowledge/topics/train_inference_consistency.md index 85c8882e..3ea8a776 100644 --- a/.agents/knowledge/topics/train_inference_consistency.md +++ b/.agents/knowledge/topics/train_inference_consistency.md @@ -36,7 +36,7 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch 4. **Model weight change**: EMA swap without restore between `inference()` and first `forward()` call. 5. **Scheduler state mismatch**: `step_index` not matching (e.g., dual-scheduler models). 6. **`num_inference_steps` changed**: invalidates sigma schedule, all trajectory timesteps are wrong. -7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is the optimize-time sample shuffle reordering `samples` before chunking into micro-batches. +7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is an optimize-time shuffle that reorders individual samples before chunking into micro-batches; adapters declare the dependence with `requires_preserved_replay_batch_composition`, and the optimize loop then shuffles whole rollout packs (see below). 8. **Stochastic conditioning encoder recomputed per forward**: When `forward()` *rebuilds* its conditioning from raw inputs each call (rather than replaying a stored embedding) and that encoder is non-deterministic, rollout and training diverge. Bagel I2I is the canonical case: the condition-image VAE (`DiagonalGaussian`, default `sample=True` -> `mean + std*randn`) is encoded **once** at rollout but **re-encoded every training `forward()`** (`_forward_rebuild` -> `_build_gen_context`), so each draws different noise -> different KV context -> on-policy `ratio != 1` (~2e-4). T2I is immune (text-only context, deterministic). Fix: make the condition encode deterministic (`vae.reg.sample = False`, posterior mean) or cache the rollout encoding and replay it. (Only affects `vae.encode` of conditions; generation uses init noise and `vae.decode` is unaffected.) **Bagel applies this fix**: `pipeline.vae.reg.sample = False` in `BagelAdapter.load_pipeline()`. 9. **Preference arms replayed with the wrong batch**: DPO shares forward-process noise across chosen/rejected states, but each arm still owns its conditioning batch. Both policy and reference forwards for the rejected state must receive `rejected_batch`; pairing it with `chosen_batch` evaluates the rejected trajectory under another conditioning context. 10. **Condition realization redrawn per candidate**: Geometry-bound VAE conditions and stochastic @@ -48,15 +48,18 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch boundary so mask/layout binding cannot drift between candidates. 11. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer. 12. **Different forward graph for the same inputs**: A prefix KV cache (prefill + cached decode) and a joint forward are mathematically equal but not bf16-equal. On Qwen-Image 2.1 they differ by 0.9–2.5% relative L2 in velocity (about 6e-6 in fp32), far outside `clip_range` and exact replay tolerances; grad mode alone does not change the result. Rollout and replay must build the same graph, so Qwen-Image 2.1 records `use_kv_cache` on the rollout sample and replay reads it back as a batch key. -13. **Execution path chosen from the batch's contents**: An adapter that runs a uniform batch batched and a mixed batch one sample at a time computes the same sample with different GEMM shapes whenever its rollout pack and its training micro-batch differ in that respect, and bf16 rounding then differs. FLUX.3 did this with its 80-token prompt windows: group-contiguous rollout packs repeat one prompt and ran batched, while shuffled micro-batches mix window counts and ran per sample, so about half the replays differed by up to 7e-6 in log-prob. Prompts long enough to span two windows trigger it, including some short captions once the chat template is added. FLUX.3 now runs every sample at transformer batch size 1 in both phases. +13. **Execution path chosen from the batch's contents**: An adapter that runs a uniform batch batched and a mixed batch one sample at a time computes the same sample with different GEMM shapes whenever its rollout pack and its training micro-batch differ in that respect, and bf16 rounding then differs. FLUX.3 did this with its 80-token prompt windows: group-contiguous rollout packs repeat one prompt and ran batched, while shuffled micro-batches mix window counts and ran per sample, so about half the replays differed by up to 7e-6 in log-prob. Prompts long enough to span two windows trigger it, including some short captions once the chat template is added. FLUX.3 therefore declares `requires_preserved_replay_batch_composition`, so replay keeps its rollout packs (item 7). ## Pack-composition-dependent adapters: `shuffle_samples` -`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. Set `train.shuffle_samples: false` so the optimize loop keeps that order: each training micro-batch then packs the *same* samples as its rollout pack -> deterministic bf16 -> on-policy `ratio == 1`. +`samples` reaches `optimize()` in rollout-pack order (`generate_samples` appends each `inference` pack in order; `compute_advantages(store_to_samples=True)` attaches advantages in place without reordering), so `samples[i*B:(i+1)*B]` is exactly rollout pack `i`. `train.shuffle_samples` keeps that pack structure whenever the adapter needs it: -- Requires matched sampling and training `per_device_batch_size` (so contiguous chunks reproduce the rollout packs). These are the *same* field: rollout uses `training_args.per_device_batch_size` for the train DataLoader (`data_utils/loader.py`) and `optimize()` chunks by the same value, so they cannot drift unless a separate sampling batch size is introduced. -- Default is `True` (per-inner-epoch shuffle); only pack-dependent adapters (Bagel) need `false`. The off-policy decorrelation cost is minor because the rollout order is already sampler-randomized. `BagelAdapter` warns at init if `shuffle_samples` is left `True`. -- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately). +- `true` (default) with a batch-invariant adapter: every inner epoch permutes individual samples. +- `true` with an adapter that declares `requires_preserved_replay_batch_composition` (Bagel, FLUX.3): every inner epoch permutes whole rollout packs and keeps each pack's samples and order, so each training micro-batch replays exactly one rollout pack -> deterministic bf16 -> on-policy `ratio == 1`. For these adapters `generate_samples` always records the rollout packs in an `AcquisitionManifest`, and the ordering fails fast if a pack does not match `per_device_batch_size` or the samples changed after rollout (for example a custom `generate_samples()` that does not record boundaries). +- `false`: the rollout order is kept for every adapter. +- Rollout and optimize chunk by the same `training_args.per_device_batch_size` (the train DataLoader in `data_utils/loader.py` and `optimize()`), so a replayed pack has its rollout size unless a separate sampling batch size is introduced. +- Shuffling packs instead of samples costs little decorrelation: the rollout order is already sampler-randomized, and with one optimizer step per epoch the accumulated gradient covers the same samples in any order. +- Wired via `BaseTrainer._order_samples_for_optimize(samples, inner_epoch)` (used by grpo/dppo/nft/awm/opd; dpo shuffles chosen/rejected pairs separately because it compares policy and reference on the same micro-batch rather than against rollout log-probs). Reward/optimization overlap keeps requiring `shuffle_samples: false` and validates the same packs against its reward tiles. ## Where in Code @@ -85,11 +88,11 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch - **Lesson**: Never flip the grad mode above code that decides from it whether a graph may outlive the call. A graph built under forced grad without a later backward is a leak whenever activation checkpointing saves non-tensor inputs that point back into it. - **Related Constraint**: N/A -### Run every FLUX.3 sample at transformer batch size 1 +### Replay FLUX.3 rollout packs whole - **Date**: 2026-10-18 - **Symptom**: FLUX.3 GRPO with `per_device_batch_size: 8` logged on-policy ratios that rounded to 1.0000 but were not exactly 1 (`ratio_std` up to 1e-4), while Qwen-Image 2.1 under the same trainer stayed exactly 1. An exact-replay probe found 421 of 768 replayed log-probs off by up to 7e-6 with GenEval prompts rewritten by the FLUX.3 upsampler, and 140 of 768 with PickScore captions. - **Root Cause**: `inference()` and `forward()` each chose batched execution for a batch whose prompts shared one 80-token window count and per-sample execution otherwise. Rollout packs of one repeated prompt ran at batch size 8; after the optimize-loop shuffle the same samples replayed at batch size 1, or the other way round. Keeping the rollout packs (`shuffle_samples: false`) made every replay exact, which isolated the pack composition. -- **Fix**: Both phases run each sample's transformer forward at batch size 1, so a sample's computation no longer depends on its batch. The regression test regroups rollout packs into new micro-batches and requires bit-identical log-probs and batch size 1 for every transformer call. +- **Fix**: FLUX.3 declares `requires_preserved_replay_batch_composition`, and `shuffle_samples` now shuffles a pack-dependent adapter's recorded rollout packs as whole units instead of individual samples, so every micro-batch replays one rollout pack in rollout order. FLUX.3 keeps batched execution for uniform-window packs. Running every sample alone was also exact but cost about 13% rollout and 15% optimization time at `per_device_batch_size: 8`. - **Lesson**: A GPU parity gate must include `per_device_batch_size > 1`, prompts of mixed lengths, and the default shuffle, and compare log-probs exactly rather than through the rounded ratio metric. Gates inherited from a `per_device_batch_size: 1` example never ran the batched path. - **Related Constraint**: Constraint #7 diff --git a/guidance/algorithms.md b/guidance/algorithms.md index 96542d4f..ed1155fb 100644 --- a/guidance/algorithms.md +++ b/guidance/algorithms.md @@ -172,7 +172,7 @@ The `kl_beta` parameter controls the coefficient of the KL divergence term. - **Full-Parameter Fine-Tuning**: The overhead is significant. You may want to set `ref_param_device` to `cpu` to save memory. - **No KL-Loss**: Setting `kl_beta` to `0` automatically disables this term and eliminates extra memory usage. -**Pack-composition-dependent adapters (e.g. Bagel)**: GRPO's on-policy step assumes `forward()` is bit-identical between rollout and training. Adapters that batch via sequence packing (NaViT) are not batch-invariant, so the default optimize-time sample shuffle breaks the on-policy `ratio == 1`. Set `train.shuffle_samples: false` (with matched sampling/training `per_device_batch_size`) for these adapters. See `workflow.md` Stage 6 and the train-inference consistency topic doc. +**Pack-composition-dependent adapters (e.g. Bagel, FLUX.3)**: GRPO's on-policy step assumes `forward()` is bit-identical between rollout and training. Adapters whose batched forward depends on which samples share a batch (NaViT sequence packing, or FLUX.3 batching only prompts that share an 80-token window count) are not batch-invariant. They declare `requires_preserved_replay_batch_composition`, and `train.shuffle_samples` then shuffles whole rollout packs instead of individual samples, so the on-policy `ratio == 1` holds under the default shuffle. See `workflow.md` Stage 6 and the train-inference consistency topic doc. #### GRPO-Guard diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md index 863ac848..14ebf98b 100644 --- a/guidance/gpu_validation.md +++ b/guidance/gpu_validation.md @@ -269,7 +269,7 @@ recorded in the resolved YAML; do not silently return to a large quality recipe. | Mode | Starting recipe | Train geometry | Steps | Algorithm-specific notes | |---|---|---|---|---| | `sd35-t2i` | `examples/grpo/lora/sd3_5/default.yaml` | `resolution: 256` | `num_inference_steps: 2` | Use the checked-in SD3.5 SFT/offline-DPO/TDM recipe as the algorithm overlay. | -| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references; keep `shuffle_samples: false`. | +| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references. | | `wan-t2v` | `examples/grpo/lora/wan21/t2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Target video has at least five frames and carries its source `fps`. | | `wan-i2v-first` | `examples/grpo/lora/wan22/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use TI2V-5B and exactly one condition image. | | `wan-flf2v` | `examples/grpo/lora/wan21/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use the dedicated FLF2V checkpoint with two ordered condition images. | diff --git a/guidance/workflow.md b/guidance/workflow.md index cd9a2df9..d6c641cb 100644 --- a/guidance/workflow.md +++ b/guidance/workflow.md @@ -625,7 +625,7 @@ For a pack-composition-dependent adapter such as Bagel at `per_device_batch_size whole work units, but validation rejects any unit that would split or repack one of those recorded batches. -> **`shuffle_samples` and on-policy ratio**: the optimize loop reorders `samples` each inner epoch (`train.shuffle_samples: true`, the default). For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT packing), this makes a training micro-batch pack a different sample set than its rollout pack, so the on-policy `ratio != 1`. Set `train.shuffle_samples: false` for such adapters (with matched sampling/training `per_device_batch_size`) so each micro-batch reproduces its rollout pack. See the train-inference consistency topic doc. +> **`shuffle_samples` and on-policy ratio**: the optimize loop reorders `samples` each inner epoch (`train.shuffle_samples: true`, the default). Batch-invariant adapters shuffle individual samples. Adapters whose batched `forward()` is *pack-composition-dependent* (Bagel NaViT packing, FLUX.3 batching only prompts that share an 80-token window count) declare `requires_preserved_replay_batch_composition`; for them the loop shuffles whole rollout packs, keeping each pack's samples and order, so every micro-batch replays one rollout pack and the on-policy `ratio == 1`. `shuffle_samples: false` keeps the rollout order for every adapter. See the train-inference consistency topic doc. ### Algorithm-Specific Optimization diff --git a/src/flow_factory/hparams/training_args/_base.py b/src/flow_factory/hparams/training_args/_base.py index 360ff6ef..a078fc59 100644 --- a/src/flow_factory/hparams/training_args/_base.py +++ b/src/flow_factory/hparams/training_args/_base.py @@ -320,12 +320,12 @@ def get_sampler_selection_contract(self) -> SamplerSelectionContract: default=True, metadata={ "help": ( - "Shuffle samples before each inner optimize epoch. Keep True normally. " - "Set False for adapters whose batched forward is pack-composition-dependent " - "(e.g. Bagel/NaViT sequence packing): then each training micro-batch packs the " - "same samples as its rollout pack, so the bf16 forward is bit-identical between " - "rollout and training and the on-policy ratio stays 1. Requires matched sampling " - "and training `per_device_batch_size`." + "Shuffle before each inner optimize epoch. Batch-invariant adapters shuffle " + "individual samples. Adapters whose batched forward depends on batch composition " + "(they declare `requires_preserved_replay_batch_composition`, e.g. Bagel and " + "FLUX.3) shuffle whole rollout packs instead, keeping each pack's samples and " + "order, so every training micro-batch replays one rollout pack and the on-policy " + "ratio stays exactly 1. False keeps the rollout order for every adapter." ) }, ) diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py index e2c69673..75d1fdcc 100644 --- a/src/flow_factory/models/abc.py +++ b/src/flow_factory/models/abc.py @@ -236,9 +236,11 @@ class BaseAdapter(ABC): supported_diffusers_cache_policies: ClassVar[Optional[frozenset[str]]] = None supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = False # Opt in when replaying a sample under a different micro-batch composition - # changes its numerical result (for example Bagel's NaViT sequence packing). - # Reward/optimization overlap then records rollout packs and refuses to tile - # or replay them with different boundaries. + # changes its numerical result, i.e. the forward is not batch-invariant (for + # example Bagel's NaViT sequence packing). The generation loop then records the + # rollout packs: `train.shuffle_samples` shuffles whole packs instead of samples, + # and reward/optimization overlap refuses to tile or replay them with different + # boundaries. requires_preserved_replay_batch_composition: ClassVar[bool] = False # Opt in only when FSDP2 communication overlap exceeds the model's activation headroom. fsdp2_use_default_stream_unshard: ClassVar[bool] = False diff --git a/src/flow_factory/models/bagel/bagel.py b/src/flow_factory/models/bagel/bagel.py index c03db026..864cc471 100644 --- a/src/flow_factory/models/bagel/bagel.py +++ b/src/flow_factory/models/bagel/bagel.py @@ -233,16 +233,6 @@ def __init__(self, config: Arguments, accelerator: Accelerator): self.pipeline: BagelPseudoPipeline self.scheduler: FlowMatchEulerDiscreteSDEScheduler - # Bagel's batched forward is pack-composition-dependent (NaViT packing), so the - # optimize-time sample shuffle must be off to keep the on-policy ratio == 1. - if getattr(self.training_args, "shuffle_samples", True): - logger.warning( - "Bagel is not batch-invariant (NaViT packing) but training_args.shuffle_samples " - "is True: this reorders optimize micro-batches vs rollout packs and breaks " - "train-inference consistency (on-policy ratio != 1). Set `train.shuffle_samples: " - "false` (see the Bagel example configs / train_inference_consistency.md)." - ) - # ─────────────────── Tokenizer & Transforms ─────────────────── def _init_tokenizer_and_transforms(self): diff --git a/src/flow_factory/models/flux/flux3.py b/src/flow_factory/models/flux/flux3.py index 474e1f4f..9de0d293 100644 --- a/src/flow_factory/models/flux/flux3.py +++ b/src/flow_factory/models/flux/flux3.py @@ -117,14 +117,17 @@ class Flux3Sample(I2ISample): class Flux3Adapter(ConfiguredImageOutputAdapterMixin, BaseAdapter): """Adapter for FLUX.3 modular text-to-image and multi-reference image-to-image models. - Rollout and replay run the transformer on one sample at a time. A sample's bf16 output then - does not depend on which samples share its rollout pack or training micro-batch, so the - on-policy ratio is exactly 1 for any ``per_device_batch_size`` and optimize-loop shuffle. + A batch whose prompts share one 80-token window count runs batched; any other batch, and + every image-conditioned batch, runs one sample at a time. A sample's bf16 result therefore + depends on the batch it is in, so the adapter declares + ``requires_preserved_replay_batch_composition`` and the optimize loop replays whole rollout + packs. """ supports_diffusers_cache: ClassVar[bool] = False supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = True ddp_find_unused_parameters: ClassVar[bool] = True + requires_preserved_replay_batch_composition: ClassVar[bool] = True component_load_dtype_defaults: ClassVar[DTypePolicy] = MappingProxyType( { "transformers": torch.bfloat16, @@ -150,6 +153,8 @@ def __init__(self, config: Arguments, accelerator: Accelerator): super().__init__(config, accelerator) self.pipeline: Flux3ModularPipeline self.scheduler: FlowMatchEulerDiscreteSDEScheduler + self._has_warned_inference_fallback = False + self._has_warned_forward_fallback = False def load_pipeline(self) -> Flux3ModularPipeline: """Load the FLUX.3 modular pipeline specification from a local path or Hub repo.""" @@ -839,7 +844,7 @@ def _inference( extra_call_back_kwargs: List[str] = [], trajectory_indices: TrajectoryIndicesType = "all", ) -> List[Flux3Sample]: - """Run FLUX.3 rollout for a single sample.""" + """Run FLUX.3 rollout on a uniform-window T2I batch or a single sample.""" device = self.device transformer = self._require_component("transformer") dtype = transformer.dtype @@ -893,19 +898,14 @@ def _inference( or not isinstance(text_ids, torch.Tensor) ): raise ValueError( - "Flux3Adapter._inference requires the stacked prompt tensors of one sample; " - "use Flux3Adapter.inference() for batches." + "Flux3Adapter._inference requires uniform stacked prompt tensors; " + "use Flux3Adapter.inference() for variable-window batches." ) prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) if isinstance(negative_prompt_embeds, torch.Tensor): negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) batch_size = prompt_embeds.shape[0] - if batch_size != 1: - raise ValueError( - f"Flux3Adapter._inference runs one sample, received batch size {batch_size}; " - "use Flux3Adapter.inference() for batches." - ) # 2. Encode condition images images = ( @@ -1113,25 +1113,82 @@ def inference( extra_call_back_kwargs: List[str] = [], trajectory_indices: TrajectoryIndicesType = "all", ) -> List[Flux3Sample]: - """Run FLUX.3 inference one sample at a time. - - Each sample's transformer forward has batch size 1, the same graph ``forward()`` replays - regardless of how training regroups samples into micro-batches. A batched forward would - round differently in bf16 whenever the rollout pack and the training micro-batch differ, - which they do under the optimize-loop shuffle or when prompts span different numbers of - 80-token windows. - """ + """Run FLUX.3 inference; falls back to per-sample execution for I2I or ragged windows.""" if isinstance(prompt, str): prompt = [prompt] if isinstance(negative_prompt, str): negative_prompt = [negative_prompt] * (len(prompt) if prompt is not None else 1) - # A flat image list or unbatched latents hold the references of a single sample. - if images is not None and not self._is_multi_images_batch(images): - images = [images] - if isinstance(image_latents, torch.Tensor) and image_latents.ndim == 2: - image_latents = image_latents.unsqueeze(0) - if isinstance(image_latent_ids, torch.Tensor) and image_latent_ids.ndim == 2: - image_latent_ids = image_latent_ids.unsqueeze(0) + + is_nested_images = self._is_multi_images_batch(images) + is_nested_image_latents = self._is_multi_image_latents(image_latents) + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + # If raw prompts are passed with batch_size > 1 and no pre-encoded embeds, check whether + # their 80-token windows differ by encoding once upfront. + if ( + prompt is not None + and len(prompt) > 1 + and prompt_embeds is None + and not (is_nested_images or is_nested_image_latents) + ): + encoded = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + guidance_scale=guidance_scale, + device=self.device, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + ) + prompt_ids = encoded["prompt_ids"] + prompt_embeds = encoded["prompt_embeds"] + text_ids = encoded["text_ids"] + negative_prompt_ids = encoded.get("negative_prompt_ids") + negative_prompt_embeds = encoded.get("negative_prompt_embeds") + negative_text_ids = encoded.get("negative_text_ids") + is_ragged_prompts = isinstance(prompt_embeds, list) or isinstance( + negative_prompt_embeds, list + ) + + fall_back = is_nested_images or is_nested_image_latents or is_ragged_prompts + if not fall_back: + return self._inference( + images=images, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + condition_image_size=condition_image_size, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + pad_multiple=pad_multiple, + prompt_ids=prompt_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + negative_prompt_ids=negative_prompt_ids, # type: ignore[arg-type] + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + condition_images=condition_images, + image_latents=image_latents, # type: ignore[arg-type] + image_latent_ids=image_latent_ids, # type: ignore[arg-type] + compute_log_prob=compute_log_prob, + extra_call_back_kwargs=extra_call_back_kwargs, + trajectory_indices=trajectory_indices, + ) + + if not self._has_warned_inference_fallback: + logger.warning( + "FLUX.3 per-sample inference fallback activated (image conditioning or " + "variable 80-token prompt windows). This warning will only appear once." + ) + self._has_warned_inference_fallback = True if images is not None: batch_size = len(images) @@ -1395,11 +1452,47 @@ def forward( "log_prob", ), ) -> FlowMatchEulerDiscreteSDESchedulerOutput: - """Replay one step for T2I and multi-reference I2I, one sample at a time. + """General forward pass supporting T2I, multi-reference I2I, and ragged prompt windows.""" + has_image_cond = image_latents is not None and ( + isinstance(image_latents, torch.Tensor) + or any(item is not None for item in image_latents) + ) + has_ragged_prompts = ( + isinstance(prompt_embeds, list) + or isinstance(text_ids, list) + or isinstance(negative_prompt_embeds, list) + or isinstance(negative_text_ids, list) + ) + fall_back = has_image_cond or has_ragged_prompts + + if not fall_back: + return self._forward( + t=t, + latents=latents, + latent_ids=latent_ids, # type: ignore[arg-type] + prompt_embeds=prompt_embeds, # type: ignore[arg-type] + text_ids=text_ids, # type: ignore[arg-type] + image_latents=None, + image_latent_ids=None, + negative_prompt_embeds=negative_prompt_embeds, # type: ignore[arg-type] + negative_text_ids=negative_text_ids, # type: ignore[arg-type] + guidance_scale=guidance_scale, + t_next=t_next, + next_latents=next_latents, + noise_level=noise_level, + attention_kwargs=attention_kwargs, + joint_attention_kwargs=joint_attention_kwargs, + compute_log_prob=compute_log_prob, + return_kwargs=return_kwargs, + ) + + if not self._has_warned_forward_fallback: + logger.warning( + "FLUX.3: Image conditioning or variable prompt window detected in forward(). " + "Processing samples individually (warning shown once)." + ) + self._has_warned_forward_fallback = True - Every sample runs through ``_forward`` with batch size 1, as in ``inference()``, so the - replayed log-probs match the rollout exactly whatever the micro-batch composition. - """ batch_size = latents.shape[0] outputs: List[FlowMatchEulerDiscreteSDESchedulerOutput] = [] diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py index a9386dda..fd7281ed 100644 --- a/src/flow_factory/trainers/abc.py +++ b/src/flow_factory/trainers/abc.py @@ -181,6 +181,48 @@ def validate_samples(self, samples: Sequence[BaseSample], *, context: str) -> No ) +def _require_rollout_packs( + manifest: Optional[AcquisitionManifest], + samples: List[BaseSample], + per_device_batch_size: int, + purpose: str, +) -> Tuple[Tuple[int, ...], ...]: + """Return the recorded rollout packs of ``samples`` for a pack-composition-dependent adapter. + + Args: + manifest: Acquisition manifest recorded by the generation loop, if any. + samples: Rank-local samples, which must still be in rollout order. + per_device_batch_size: Optimize micro-batch size every rollout pack must match. + purpose: What needs the packs; it prefixes the missing-metadata error. + + Returns: + One tuple of sample indices per rollout micro-batch, in rollout order. + + Raises: + RuntimeError: If the generation loop recorded no rollout micro-batches. + ValueError: If a rollout micro-batch does not match ``per_device_batch_size`` or the + samples changed since rollout. + """ + if manifest is None or not manifest.rollout_batches: + raise RuntimeError( + f"{purpose} for a pack-composition-dependent adapter requires rollout " + "micro-batch metadata; custom generate_samples() implementations must " + "record the base generation-loop batch boundaries" + ) + + invalid_sizes = tuple( + len(batch) for batch in manifest.rollout_batches if len(batch) != per_device_batch_size + ) + if invalid_sizes: + raise ValueError( + "pack-composition-dependent replay requires every rollout micro-batch to match " + f"train.per_device_batch_size={per_device_batch_size}; invalid_sizes={invalid_sizes!r}" + ) + + manifest.validate_samples(samples) + return manifest.rollout_batches + + class BaseTrainer(MultiRoleCheckpointingMixin, MultiRoleBackendValidationMixin, ABC): """Abstract base class for Flow-Factory trainers.""" @@ -231,6 +273,7 @@ def __init__( self._acquisition_cycle_active = False self._acquisition_cycle_incomplete = False self._reward_overlap_acquisition_manifest: Optional[AcquisitionManifest] = None + self._rollout_acquisition_manifest: Optional[AcquisitionManifest] = None self.acquisition_driver: AcquisitionDriver = build_acquisition_driver( type(self).execution_contract ) @@ -2241,24 +2284,8 @@ def _validate_reward_overlap_replay_batch_composition( ) manifest = plan.manifest - if manifest is None or not manifest.rollout_batches: - raise RuntimeError( - "reward overlap for a pack-composition-dependent adapter requires rollout " - "micro-batch metadata; custom generate_samples() implementations must " - "record the base generation-loop batch boundaries" - ) - batch_size = self.training_args.per_device_batch_size - invalid_sizes = tuple( - len(batch) for batch in manifest.rollout_batches if len(batch) != batch_size - ) - if invalid_sizes: - raise ValueError( - "pack-composition-dependent replay requires every rollout micro-batch to " - f"match train.per_device_batch_size={batch_size}; invalid_sizes={invalid_sizes!r}" - ) - - manifest.validate_samples(samples) + _require_rollout_packs(manifest, samples, batch_size, "reward overlap") boundaries = manifest.rollout_boundaries for tile in plan.tiles: if tile.start not in boundaries or tile.stop not in boundaries: @@ -2865,17 +2892,29 @@ def _order_samples_for_optimize( ) -> List[BaseSample]: """Return the per-inner-epoch sample ordering for the optimize loop. - When ``training_args.shuffle_samples`` is False, the rollout-pack order is - preserved so each training micro-batch packs exactly the samples of its - corresponding rollout ``inference`` pack. For adapters whose batched forward - is pack-composition-dependent (e.g. Bagel/NaViT packing), this keeps the - bf16 forward bit-identical between rollout and training (on-policy ratio==1). + When ``training_args.shuffle_samples`` is False, the rollout order is kept, so each + training micro-batch packs exactly the samples of its rollout ``inference`` pack. + Otherwise a batch-invariant adapter's samples are shuffled individually. An adapter + that declares ``requires_preserved_replay_batch_composition`` (its batched forward + depends on which samples share a batch, e.g. Bagel's NaViT packing or FLUX.3's + window-dependent batching) has its recorded rollout packs shuffled as whole units: + every micro-batch then replays one rollout pack with its samples in rollout order, + so the bf16 forward stays bit-identical to the rollout (on-policy ratio == 1). """ if not self.training_args.shuffle_samples: return samples perm_gen = create_generator(self.training_args.seed, self.epoch, inner_epoch) - perm = torch.randperm(len(samples), generator=perm_gen) - return [samples[i] for i in perm] + if not self.adapter.requires_preserved_replay_batch_composition: + perm = torch.randperm(len(samples), generator=perm_gen) + return [samples[i] for i in perm] + packs = _require_rollout_packs( + self._rollout_acquisition_manifest, + samples, + self.training_args.per_device_batch_size, + "optimize-loop pack shuffling", + ) + pack_perm = torch.randperm(len(packs), generator=perm_gen) + return [samples[i] for pack_index in pack_perm.tolist() for i in packs[pack_index]] def _maybe_offload_samples_to_cpu(self, samples: List[BaseSample]) -> None: """Offload each sample's tensors to pinned CPU when offload is enabled. @@ -3161,8 +3200,12 @@ def generate_samples( reward_buffer is not None and getattr(self.training_args, "reward_optimization_overlap", False) ) + record_rollout_batches = ( + capture_overlap_manifest or self.adapter.requires_preserved_replay_batch_composition + ) recorded_rollout_batches: List[Tuple[int, ...]] = [] self._reward_overlap_acquisition_manifest = None + self._rollout_acquisition_manifest = None if reward_buffer is not None: reward_buffer.clear() @@ -3196,15 +3239,17 @@ def generate_samples( trajectory_indices=trajectory_indices, **extra_inference_kwargs, ) - if capture_overlap_manifest: + if record_rollout_batches: recorded_rollout_batches.append(tuple(id(sample) for sample in sample_batch)) samples.extend(sample_batch) - if capture_overlap_manifest: - self._reward_overlap_acquisition_manifest = AcquisitionManifest.from_samples( + if record_rollout_batches: + self._rollout_acquisition_manifest = AcquisitionManifest.from_samples( samples, rollout_batch_object_ids=recorded_rollout_batches, ) + if capture_overlap_manifest: + self._reward_overlap_acquisition_manifest = self._rollout_acquisition_manifest # Multi-source invariant: when more than one training source is # active, batches flow through `MultiSourceTrainDataLoader`, which diff --git a/tests/models/test_flux3.py b/tests/models/test_flux3.py index 508d7111..7828b676 100644 --- a/tests/models/test_flux3.py +++ b/tests/models/test_flux3.py @@ -287,6 +287,8 @@ def _build_flux3_adapter(*, flow_map: bool = False, noise_level: float = 0.0) -> ) adapter._output_state_codec = adapter._build_output_state_codec_declaration() adapter._output_state_encoding_modules = adapter._validate_output_state_codec_lifecycle() + adapter._has_warned_inference_fallback = False + adapter._has_warned_forward_fallback = False return adapter @@ -619,7 +621,10 @@ def test_sde_rollout_replay_parity_and_gradient_flow() -> None: assert adapter.transformer.single_transformer_blocks[0].attn.to_q.weight.grad is not None -def test_on_policy_replay_is_exact_for_any_micro_batch_composition() -> None: +def test_rollout_packs_replay_exactly_through_their_rollout_path() -> None: + # Window-dependent batching makes the forward pack-dependent; the optimize loop then keeps + # every rollout pack intact, so replaying a pack must reproduce its rollout bit for bit. + assert Flux3Adapter.requires_preserved_replay_batch_composition is True adapter = _build_flux3_adapter(flow_map=False, noise_level=0.7) adapter.transformer.train() transformer_batch_sizes: list[int] = [] @@ -630,29 +635,22 @@ def test_on_policy_replay_is_exact_for_any_micro_batch_composition() -> None: with_kwargs=True, ) short_prompt, long_prompt = "fox", "word " * 15 # one and two 80-token windows - rollout_kwargs = dict( - height=16, - width=16, - num_inference_steps=3, - guidance_scale=2.0, - text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), - compute_log_prob=True, - ) - # A uniform pack, as group-contiguous sampling produces, and a pack of mixed window counts. - samples = adapter.inference( - prompt=[short_prompt, short_prompt], - generator=torch.Generator().manual_seed(7), - **rollout_kwargs, - ) + adapter.inference( - prompt=[short_prompt, long_prompt], - generator=torch.Generator().manual_seed(8), - **rollout_kwargs, - ) - assert set(transformer_batch_sizes) == {1} + for pack, batch_size in (([short_prompt, short_prompt], 2), ([short_prompt, long_prompt], 1)): + transformer_batch_sizes.clear() + samples = adapter.inference( + prompt=pack, + height=16, + width=16, + num_inference_steps=3, + guidance_scale=2.0, + text_encoder_out_layers=(4, 8, 12, 16, 20, 24, 28, 32), + generator=torch.Generator().manual_seed(7), + compute_log_prob=True, + ) + assert set(transformer_batch_sizes) == {batch_size} - # Optimize micro-batches regroup the rollout packs, as the optimize-loop shuffle does. - for order in ([0, 2], [1, 3], [3, 2, 1, 0]): - batch = BaseSample.stack([samples[index] for index in order]) + transformer_batch_sizes.clear() + batch = BaseSample.stack(samples) for step_idx in range(2): t = batch["timesteps"][:, step_idx] replay = adapter.forward( @@ -670,7 +668,7 @@ def test_on_policy_replay_is_exact_for_any_micro_batch_composition() -> None: compute_log_prob=True, ) assert torch.equal(replay.log_prob.detach(), batch["log_probs"][:, step_idx]) - assert set(transformer_batch_sizes) == {1} + assert set(transformer_batch_sizes) == {batch_size} def test_gradient_checkpointing_order_and_lora_round_trip(tmp_path: Path) -> None: diff --git a/tests/trainers/test_optimize_sample_order.py b/tests/trainers/test_optimize_sample_order.py new file mode 100644 index 00000000..57a01c14 --- /dev/null +++ b/tests/trainers/test_optimize_sample_order.py @@ -0,0 +1,181 @@ +# Copyright 2026 Jayce-Ping +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optimize-loop sample ordering for batch-invariant and pack-dependent adapters.""" + +from __future__ import annotations + +import contextlib +from types import SimpleNamespace +from typing import Any, Iterator, List, Optional + +import pytest +import torch + +from flow_factory.samples import AcquisitionManifest +from flow_factory.trainers.abc import BaseTrainer +from flow_factory.utils.base import create_generator + +SEED, EPOCH = 42, 3 + + +def _samples(num_packs: int, pack_size: int) -> List[SimpleNamespace]: + return [ + SimpleNamespace(unique_id=pack, name=f"pack{pack}-item{item}") + for pack in range(num_packs) + for item in range(pack_size) + ] + + +def _packs_manifest(samples: List[Any], pack_size: int) -> AcquisitionManifest: + return AcquisitionManifest.from_samples( + samples, + rollout_batch_object_ids=[ + [id(sample) for sample in samples[start : start + pack_size]] + for start in range(0, len(samples), pack_size) + ], + ) + + +def _trainer( + *, + pack_dependent: bool, + shuffle: bool = True, + per_device_batch_size: int = 2, + manifest: Optional[AcquisitionManifest] = None, +) -> SimpleNamespace: + return SimpleNamespace( + training_args=SimpleNamespace( + shuffle_samples=shuffle, seed=SEED, per_device_batch_size=per_device_batch_size + ), + epoch=EPOCH, + adapter=SimpleNamespace(requires_preserved_replay_batch_composition=pack_dependent), + _rollout_acquisition_manifest=manifest, + ) + + +def _names(samples: List[Any]) -> List[str]: + return [sample.name for sample in samples] + + +@pytest.mark.parametrize("pack_dependent", [False, True]) +def test_shuffle_samples_false_keeps_the_rollout_order(pack_dependent: bool) -> None: + samples = _samples(num_packs=4, pack_size=2) + trainer = _trainer(pack_dependent=pack_dependent, shuffle=False) + + assert BaseTrainer._order_samples_for_optimize(trainer, samples, 0) == samples + + +def test_batch_invariant_adapters_shuffle_individual_samples() -> None: + samples = _samples(num_packs=6, pack_size=2) + trainer = _trainer(pack_dependent=False) + + ordered = BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) + + permutation = torch.randperm(len(samples), generator=create_generator(SEED, EPOCH, 1)) + assert ordered == [samples[index] for index in permutation] + assert ordered != samples + + +def test_pack_dependent_adapters_shuffle_whole_rollout_packs() -> None: + pack_size = 2 + samples = _samples(num_packs=6, pack_size=pack_size) + trainer = _trainer(pack_dependent=True, manifest=_packs_manifest(samples, pack_size)) + + ordered = BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) + + micro_batches = [ + _names(ordered[start : start + pack_size]) for start in range(0, len(ordered), pack_size) + ] + rollout_packs = [ + _names(samples[start : start + pack_size]) for start in range(0, len(samples), pack_size) + ] + pack_permutation = torch.randperm( + len(rollout_packs), generator=create_generator(SEED, EPOCH, 1) + ) + assert micro_batches == [rollout_packs[index] for index in pack_permutation] + assert micro_batches != rollout_packs + assert BaseTrainer._order_samples_for_optimize(trainer, samples, inner_epoch=1) == ordered + + +def test_single_sample_packs_shuffle_like_individual_samples() -> None: + samples = _samples(num_packs=8, pack_size=1) + pack_dependent = _trainer( + pack_dependent=True, per_device_batch_size=1, manifest=_packs_manifest(samples, 1) + ) + batch_invariant = _trainer(pack_dependent=False, per_device_batch_size=1) + + assert BaseTrainer._order_samples_for_optimize( + pack_dependent, samples, 2 + ) == BaseTrainer._order_samples_for_optimize(batch_invariant, samples, 2) + + +def test_pack_shuffling_requires_matching_rollout_packs() -> None: + samples = _samples(num_packs=3, pack_size=2) + + with pytest.raises(RuntimeError, match="optimize-loop pack shuffling .* rollout micro-batch"): + BaseTrainer._order_samples_for_optimize(_trainer(pack_dependent=True), samples, 0) + + uneven = _trainer( + pack_dependent=True, + per_device_batch_size=2, + manifest=_packs_manifest(samples, pack_size=3), + ) + with pytest.raises(ValueError, match=r"train.per_device_batch_size=2; invalid_sizes=\(3, 3\)"): + BaseTrainer._order_samples_for_optimize(uneven, samples, 0) + + reordered = _trainer(pack_dependent=True, manifest=_packs_manifest(samples, pack_size=2)) + with pytest.raises(ValueError, match="samples changed after its manifest was created"): + BaseTrainer._order_samples_for_optimize(reordered, samples[::-1], 0) + + +class _GenerationHarness: + def __init__(self, pack_dependent: bool) -> None: + self.adapter = SimpleNamespace( + rollout=lambda: None, requires_preserved_replay_batch_composition=pack_dependent + ) + self.training_args = SimpleNamespace( + reward_optimization_overlap=False, num_batches_per_epoch=2 + ) + self.dataloader = [{"pack": 0}, {"pack": 1}] + self.epoch = 0 + self.show_progress_bar = False + self.train_dataloaders_by_source = {"default": self.dataloader} + + @contextlib.contextmanager + def _rollout_acceleration(self) -> Iterator[None]: + yield + + @contextlib.contextmanager + def autocast(self) -> Iterator[None]: + yield + + def sample_batch(self, batch: dict, **_: Any) -> List[SimpleNamespace]: + return [SimpleNamespace(unique_id=batch["pack"], source=None) for _ in range(2)] + + +@pytest.mark.parametrize("pack_dependent", [False, True]) +def test_generation_records_rollout_packs_for_pack_dependent_adapters( + pack_dependent: bool, +) -> None: + harness = _GenerationHarness(pack_dependent) + + samples = BaseTrainer.generate_samples(harness) + + assert harness._reward_overlap_acquisition_manifest is None + if pack_dependent: + assert harness._rollout_acquisition_manifest.rollout_batches == ((0, 1), (2, 3)) + harness._rollout_acquisition_manifest.validate_samples(samples) + else: + assert harness._rollout_acquisition_manifest is None