`flux2` model/pipeline review

flux2 model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Test execution note: focused reproductions were run with .venv. Full Flux2 pytest collection was attempted, but collection fails before the Flux2 tests run on this local ROCm Windows torch build. The failure path is the shared test import chain into diffusers.training_utils, which imports torch.distributed.fsdp. This torch build exposes the torch.distributed module object, but torch.distributed.is_available() is False; importing FSDP then requires torch._C._distributed_c10d and raises ModuleNotFoundError. The guard in src/diffusers/training_utils.py currently checks only getattr(torch, "distributed", None) is not None, which is not sufficient for torch builds without distributed support. It should use torch.distributed.is_available() as in src/diffusers/models/attention_dispatch.py.

Duplicate search status: searched existing GitHub Issues and PRs for Flux2, affected class/function/file names, and the specific failure modes below. No likely duplicates were found.

Issue 1: Klein pipelines return decoded-shape tensors for output_type="latent"

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py#L903-L914
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py#L866-L878
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py#L1244-L1257

Problem:
Flux2KleinPipeline, Flux2KleinKVPipeline, and Flux2KleinInpaintPipeline unpack, denormalize, and unpatchify latents before checking output_type == "latent". By contrast, Flux2Pipeline returns the packed transformer latents directly before VAE decode/unpack work.

Impact:
Users requesting latent output from Klein variants receive a different tensor layout from the base Flux2 pipeline. This breaks pipeline interchangeability and downstream latent workflows that expect the packed latent shape.

Reproduction:

import torch
from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2KleinPipeline, Flux2Pipeline, Flux2Transformer2DModel


def tiny_vae():
    return AutoencoderKLFlux2(
        in_channels=3,
        out_channels=3,
        down_block_types=("DownEncoderBlock2D",),
        up_block_types=("UpDecoderBlock2D",),
        block_out_channels=(4,),
        layers_per_block=1,
        latent_channels=1,
        norm_num_groups=1,
        sample_size=4,
        mid_block_add_attention=False,
    )


def tiny_transformer(guidance_embeds):
    return Flux2Transformer2DModel(
        in_channels=4,
        out_channels=4,
        num_layers=0,
        num_single_layers=0,
        attention_head_dim=8,
        num_attention_heads=1,
        joint_attention_dim=8,
        timestep_guidance_channels=8,
        axes_dims_rope=(2, 2, 2, 2),
        guidance_embeds=guidance_embeds,
    )


prompt_embeds = torch.randn(1, 2, 8)

std = Flux2Pipeline(FlowMatchEulerDiscreteScheduler(), tiny_vae(), None, None, tiny_transformer(True))
klein = Flux2KleinPipeline(
    FlowMatchEulerDiscreteScheduler(), tiny_vae(), None, None, tiny_transformer(False), is_distilled=True
)

print(tuple(std(prompt_embeds=prompt_embeds, height=4, width=4, num_inference_steps=1, output_type="latent").images.shape))
print(tuple(klein(prompt_embeds=prompt_embeds, height=4, width=4, num_inference_steps=1, output_type="latent").images.shape))
# Current output: (1, 4, 4) vs (1, 1, 4, 4)

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2.py#L1011-L1023

Suggested fix:
Move the latent-output branch before unpack/denormalize/unpatchify in all three Klein variants.

if output_type == "latent":
    image = latents
else:
    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
    latents = self._denormalize_latents(latents, self.vae, self.vae_scale_factor)
    latents = self.image_processor.unpatchify(latents)
    image = self.vae.decode(latents, return_dict=False)[0]
    image = self.image_processor.postprocess(image, output_type=output_type)

Issue 2: Modular Flux2 decode step decodes even for output_type="latent"

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/decoders.py#L166-L180

Problem:
Flux2DecodeStep always unpacks, denormalizes, unpatchifies, and calls vae.decode() before postprocessing. It does not branch around decode work when output_type == "latent".

Impact:
Modular Flux2 cannot return true latent outputs without requiring a VAE decode path. This is inconsistent with the non-modular pipeline contract and can make latent-only workflows slower or fail when the VAE should not be needed.

Reproduction:

import torch
from diffusers.modular_pipelines.flux2.decoders import Flux2DecodeStep
from diffusers.modular_pipelines.modular_pipeline import PipelineState


class BN:
    running_mean = torch.zeros(4)
    running_var = torch.ones(4)


class Config:
    batch_norm_eps = 1e-4


class VAE:
    bn = BN()
    config = Config()

    def decode(self, *args, **kwargs):
        raise RuntimeError("decode() should not be called for output_type='latent'")


class ImageProcessor:
    def postprocess(self, image, output_type):
        return image


class Components:
    vae = VAE()
    image_processor = ImageProcessor()


state = PipelineState()
state.set("latents", torch.zeros(1, 4, 2, 2))
state.set("output_type", "latent")

Flux2DecodeStep()(Components(), state)

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2.py#L1011-L1023

Suggested fix:
Add the same early latent-output branch used by Flux2Pipeline before VAE-specific decode work.

if block_state.output_type == "latent":
    block_state.images = block_state.latents
else:
    block_state.latents = self._unpack_latents(
        block_state.latents,
        block_state.height,
        block_state.width,
        components.vae_scale_factor,
    )
    block_state.latents = self._denormalize_latents(
        block_state.latents,
        components.vae,
        components.vae_scale_factor,
    )
    block_state.latents = components.image_processor.unpatchify(block_state.latents)
    block_state.images = components.vae.decode(block_state.latents, return_dict=False)[0]
    block_state.images = components.image_processor.postprocess(
        block_state.images, output_type=block_state.output_type
    )

Issue 3: Modular Klein base CFG path does not use transformer cache contexts

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/denoise.py#L338-L356

Problem:
Flux2KleinBaseLoopDenoiser loops over guider batches and calls the transformer without components.transformer.cache_context(...). The monolithic Klein pipeline wraps each conditional/unconditional transformer call in a cache context keyed by the guider batch identifier.

Impact:
The modular Klein base path can lose the intended per-branch cache separation and reuse behavior. That is especially risky for CFG, where conditional and unconditional transformer calls must not accidentally share the wrong cache state.

Reproduction:

from contextlib import contextmanager
import torch
from diffusers.modular_pipelines.modular_pipeline import BlockState
from diffusers.modular_pipelines.flux2.denoise import Flux2KleinBaseLoopDenoiser


class Batch:
    def __init__(self, name, prompt, txt_ids):
        self.name = name
        self.encoder_hidden_states = prompt
        self.txt_ids = txt_ids


class Guider:
    _identifier_key = "name"

    def set_state(self, **kwargs):
        pass

    def prepare_inputs(self, inputs):
        return [
            Batch("cond", inputs["encoder_hidden_states"][0], inputs["txt_ids"][0]),
            Batch("uncond", inputs["encoder_hidden_states"][1], inputs["txt_ids"][1]),
        ]

    def prepare_models(self, model):
        pass

    def cleanup_models(self, model):
        pass

    def __call__(self, state):
        return (state[0].noise_pred,)


class Transformer:
    dtype = torch.float32
    active_context = None
    seen_contexts = []

    @contextmanager
    def cache_context(self, name):
        self.active_context = name
        yield
        self.active_context = None

    def __call__(self, hidden_states, **kwargs):
        self.seen_contexts.append(self.active_context)
        return (torch.zeros_like(hidden_states),)


class Components:
    transformer = Transformer()
    guider = Guider()


block_state = BlockState(
    latents=torch.zeros(1, 2, 4),
    latent_ids=torch.zeros(1, 2, 4),
    image_latents=None,
    prompt_embeds=torch.zeros(1, 3, 4),
    negative_prompt_embeds=torch.zeros(1, 3, 4),
    txt_ids=torch.zeros(1, 3, 4),
    negative_txt_ids=torch.zeros(1, 3, 4),
    joint_attention_kwargs=None,
    num_inference_steps=1,
)
Flux2KleinBaseLoopDenoiser()(Components(), block_state, 0, torch.tensor(1.0))
print(Components.transformer.seen_contexts)
# Current output: [None, None]

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py#L844-L866

Suggested fix:
Wrap the modular transformer call in the same cache context used by the monolithic Klein pipeline.

context_name = getattr(guider_state_batch, components.guider._identifier_key)
with components.transformer.cache_context(context_name):
    noise_pred = components.transformer(
        hidden_states=block_state.latents,
        timestep=timestep / 1000,
        guidance=None,
        encoder_hidden_states=guider_state_batch.encoder_hidden_states,
        txt_ids=guider_state_batch.txt_ids,
        img_ids=block_state.latent_ids,
        joint_attention_kwargs=block_state.joint_attention_kwargs,
        return_dict=False,
    )[0]

Issue 4: Flux2KleinKVPipeline.is_distilled is not serialized

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py#L183-L210

Problem:
Flux2KleinKVPipeline.__init__ accepts is_distilled, stores it as an instance attribute, but never registers it in the pipeline config. Flux2KleinPipeline registers the same argument with register_to_config.

Impact:
Saved Flux2KleinKVPipeline configs do not preserve whether the pipeline is distilled. Reloading the pipeline can silently fall back to the constructor default instead of the original setting.

Reproduction:

from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2KleinKVPipeline, Flux2Transformer2DModel

vae = AutoencoderKLFlux2(
    in_channels=3,
    out_channels=3,
    down_block_types=("DownEncoderBlock2D",),
    up_block_types=("UpDecoderBlock2D",),
    block_out_channels=(4,),
    layers_per_block=1,
    latent_channels=1,
    norm_num_groups=1,
    sample_size=4,
    mid_block_add_attention=False,
)
transformer = Flux2Transformer2DModel(
    in_channels=4,
    out_channels=4,
    num_layers=0,
    num_single_layers=0,
    attention_head_dim=8,
    num_attention_heads=1,
    joint_attention_dim=8,
    timestep_guidance_channels=8,
    axes_dims_rope=(2, 2, 2, 2),
    guidance_embeds=False,
)
pipe = Flux2KleinKVPipeline(
    FlowMatchEulerDiscreteScheduler(), vae, None, None, transformer, is_distilled=False
)
print("is_distilled" in pipe.config, pipe.config.get("is_distilled", None))
# Current output: False None

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py#L186-L198

Suggested fix:
Register is_distilled in the KV pipeline constructor.

self.register_to_config(is_distilled=is_distilled)

Issue 5: Flux2 modular blocks depend on non-modular pipeline modules and have generated TODO docs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/inputs.py#L18
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/decoders.py#L25
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py#L285-L332
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py#L324-L372
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py#L337-L385

Problem:
Flux2 modular blocks import Flux2ImageProcessor from diffusers.pipelines.flux2.image_processor, so modular pipeline code depends on the non-modular pipeline package. The generated modular block files also still contain TODO: Add description placeholders in public input/output docs.

Impact:
This violates the modular pipeline layering rules and leaves public modular pipeline documentation incomplete. It also makes future pipeline refactors more fragile because modular code is coupled to non-modular package internals.

Reproduction:

from pathlib import Path

for path in Path("src/diffusers/modular_pipelines/flux2").glob("*.py"):
    text = path.read_text()
    if "from ...pipelines.flux2" in text or "TODO: Add description" in text:
        print(path)

Relevant precedent:
Other modular pipeline families keep shared utilities outside the non-modular pipeline package or avoid importing from the non-modular pipeline implementation. Modular docs are expected to be generated without TODO placeholders.

Suggested fix:
Move Flux2ImageProcessor to a shared non-pipeline module, or introduce a modular/shared image processor location and update both modular and non-modular imports to use it. Fill in the missing InputParam and OutputParam descriptions, then regenerate with:

python utils/modular_auto_docstring.py --fix_and_overwrite

Issue 6: Flux2 coverage is missing slow tests and some public docs/model tests

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux2/test_pipeline_flux2.py#L23-L24
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux2/test_pipeline_flux2_klein.py#L19-L20
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py#L26-L27
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py#L19-L20
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py#L38-L71
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/flux2.md#L35-L49

Problem:
The Flux2 pipeline tests are fast-only. No slow Flux2 pipeline tests were found for Flux2Pipeline, Flux2KleinPipeline, Flux2KleinKVPipeline, or Flux2KleinInpaintPipeline. AutoencoderKLFlux2 is a public model class but does not appear to have a dedicated autoencoder test file, and the Flux2 docs page omits Flux2KleinInpaintPipeline. No dedicated AutoencoderKLFlux2 API docs page was found.

Impact:
The family lacks coverage that exercises real checkpoints, save/load behavior, and public API documentation for all exposed classes. This increases regression risk for pipeline parity, model serialization, and user-facing docs.

Reproduction:

from pathlib import Path

print("slow markers:", [
    str(p) for p in Path("tests").rglob("*flux2*.py")
    if "@slow" in p.read_text(errors="ignore")
])
print("autoencoder tests:", list(Path("tests/models/autoencoders").glob("*flux2*.py")))
docs = Path("docs/source/en/api/pipelines/flux2.md").read_text()
print("inpaint docs:", "Flux2KleinInpaintPipeline" in docs)

Relevant precedent:
Established model/pipeline families generally include fast tests, slow checkpoint tests, model save/load coverage, and docs entries for each public pipeline/model class.

Suggested fix:
Add slow tests for the standard, Klein, Klein KV, and Klein inpaint pipelines using the smallest public checkpoints that cover the public APIs. Add tests/models/autoencoders/test_models_autoencoder_kl_flux2.py for AutoencoderKLFlux2 config/save-load behavior. Add autodoc sections for Flux2KleinInpaintPipeline and AutoencoderKLFlux2, plus the model docs toctree entry where appropriate.

Issue 7: Shared test imports require torch.distributed even when the installed torch build does not support it

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/training_utils.py#L16-L20

Problem:
training_utils.py guards FSDP imports with getattr(torch, "distributed", None) is not None. That is weaker than checking distributed support. On torch builds that expose torch.distributed but were built without distributed backend support, such as the local ROCm Windows build used for this audit, torch.distributed.is_available() is False. Importing torch.distributed.fsdp then requires torch._C._distributed_c10d and fails during test collection.

Impact:
Flux2 test collection, and any other tests importing diffusers.training_utils, can fail before target tests run on valid torch installs without distributed support. This unnecessarily blocks local CPU/ROCm Windows validation of unrelated model and pipeline behavior.

Reproduction:

import torch

print(torch.distributed.is_available())
from diffusers import training_utils

On the affected environment this fails while importing torch.distributed.fsdp with:

ModuleNotFoundError: No module named 'torch._C._distributed_c10d'

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_dispatch.py#L26-L31

Suggested fix:
Use torch.distributed.is_available() before importing FSDP symbols.

if getattr(torch, "distributed", None) is not None and torch.distributed.is_available():
    from torch.distributed.fsdp import CPUOffload, ShardingStrategy
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
else:
    CPUOffload = None
    ShardingStrategy = None
    FSDP = None
    transformer_auto_wrap_policy = None

Start with the affected Flux2 pipeline files, modular_pipelines/flux2/decoders.py, modular_pipelines/flux2/denoise.py, and training_utils.py, comparing them with pipeline_flux2.py, pipeline_flux2_klein.py, and models/attention_dispatch.py. Run the focused reproductions under .venv, then verify latent outputs, cache contexts, serialized configuration, and collection behavior; completion requires all four reported failures to be addressed.