`flux2` model/pipeline review

[hlky]

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

Lines 903 to 914 in 0f1abc4

	latents = self._unpack_latents_with_ids(latents, latent_ids, latent_height // 2, latent_width // 2)
	latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	if output_type == "latent":
	image = latents
	else:
	image = self.vae.decode(latents, return_dict=False)[0]

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py

Lines 866 to 878 in 0f1abc4

	latents = self._unpack_latents_with_ids(latents, latent_ids)
	latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	if output_type == "latent":
	image = latents
	else:
	image = self.vae.decode(latents, return_dict=False)[0]
	image = self.image_processor.postprocess(image, output_type=output_type)

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py

Lines 1244 to 1257 in 0f1abc4

	latents = self._unpack_latents_with_ids(latents, latent_image_ids)
	latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	if output_type == "latent":
	image = latents
	else:
	image = self.vae.decode(latents, return_dict=False)[0]
	image = self.image_processor.postprocess(image, output_type=output_type)

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2.py

Lines 1011 to 1023 in 0f1abc4

	if output_type == "latent":
	image = latents
	else:
	latents = self._unpack_latents_with_ids(latents, latent_ids)
	latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	image = self.vae.decode(latents, return_dict=False)[0]

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:

diffusers/src/diffusers/modular_pipelines/flux2/decoders.py

Lines 166 to 180 in 0f1abc4

	block_state = self.get_block_state(state)
	vae = components.vae
	latents = block_state.latents
	latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	block_state.images = vae.decode(latents, return_dict=False)[0]
	block_state.images = components.image_processor.postprocess(

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2.py

Lines 1011 to 1023 in 0f1abc4

	if output_type == "latent":
	image = latents
	else:
	latents = self._unpack_latents_with_ids(latents, latent_ids)
	latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
	latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
	latents.device, latents.dtype
	)
	latents = latents * latents_bn_std + latents_bn_mean
	latents = self._unpatchify_latents(latents)
	image = self.vae.decode(latents, return_dict=False)[0]

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:

diffusers/src/diffusers/modular_pipelines/flux2/denoise.py

Lines 338 to 356 in 0f1abc4

	guider_state = components.guider.prepare_inputs(guider_inputs)
	for guider_state_batch in guider_state:
	components.guider.prepare_models(components.transformer)
	cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
	noise_pred = components.transformer(
	hidden_states=latent_model_input,
	timestep=timestep / 1000,
	guidance=None,
	img_ids=img_ids,
	joint_attention_kwargs=block_state.joint_attention_kwargs,
	return_dict=False,
	**cond_kwargs,
	)[0]
	guider_state_batch.noise_pred = noise_pred[:, : latents.size(1)]
	components.guider.cleanup_models(components.transformer)
	# perform guidance

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

Lines 844 to 866 in 0f1abc4

	with self.transformer.cache_context("cond"):
	noise_pred = self.transformer(
	hidden_states=latent_model_input, # (B, image_seq_len, C)
	timestep=timestep / 1000,
	guidance=None,
	encoder_hidden_states=prompt_embeds,
	txt_ids=text_ids, # B, text_seq_len, 4
	img_ids=latent_image_ids, # B, image_seq_len, 4
	joint_attention_kwargs=self.attention_kwargs,
	return_dict=False,
	)[0]
	noise_pred = noise_pred[:, : latents.size(1) :]
	if self.do_classifier_free_guidance:
	with self.transformer.cache_context("uncond"):
	neg_noise_pred = self.transformer(
	hidden_states=latent_model_input,
	timestep=timestep / 1000,
	guidance=None,
	encoder_hidden_states=negative_prompt_embeds,
	txt_ids=negative_text_ids,
	img_ids=latent_image_ids,

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py

Lines 183 to 210 in 0f1abc4

	def __init__(
	self,
	scheduler: FlowMatchEulerDiscreteScheduler,
	vae: AutoencoderKLFlux2,
	text_encoder: Qwen3ForCausalLM,
	tokenizer: Qwen2TokenizerFast,
	transformer: Flux2Transformer2DModel,
	is_distilled: bool = True,
	):
	super().__init__()
	self.register_modules(
	vae=vae,
	text_encoder=text_encoder,
	tokenizer=tokenizer,
	scheduler=scheduler,
	transformer=transformer,
	)
	self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
	# Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
	# by the patch size. So the vae scale factor is multiplied by the patch size to account for this
	self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
	self.tokenizer_max_length = 512
	self.default_sample_size = 128
	# Set KV-cache-aware attention processors
	self._set_kv_attn_processors()

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

Lines 186 to 198 in 0f1abc4

	is_distilled: bool = False,
	):
	super().__init__()
	self.register_modules(
	vae=vae,
	text_encoder=text_encoder,
	tokenizer=tokenizer,
	scheduler=scheduler,
	transformer=transformer,
	)
	self.register_to_config(is_distilled=is_distilled)

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:

diffusers/src/diffusers/modular_pipelines/flux2/inputs.py

Line 18 in 0f1abc4

from ...pipelines.flux2.image_processor import Flux2ImageProcessor

diffusers/src/diffusers/modular_pipelines/flux2/decoders.py

Line 25 in 0f1abc4

from ...pipelines.flux2.image_processor import Flux2ImageProcessor

diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2.py

Lines 285 to 332 in 0f1abc4

	class Flux2AutoBlocks(SequentialPipelineBlocks):
	"""
	Auto Modular pipeline for text-to-image and image-conditioned generation using Flux2.
	Supported workflows:
	- `text2image`: requires `prompt`
	- `image_conditioned`: requires `image`, `prompt`
	Components:
	text_encoder (`Mistral3ForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
	(`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
	(`Flux2Transformer2DModel`)
	Inputs:
	prompt (`None`, *optional*):
	TODO: Add description.
	max_sequence_length (`int`, *optional*, defaults to 512):
	TODO: Add description.
	text_encoder_out_layers (`tuple`, *optional*, defaults to (10, 20, 30)):
	TODO: Add description.
	image (`None`, *optional*):
	TODO: Add description.
	height (`None`, *optional*):
	TODO: Add description.
	width (`None`, *optional*):
	TODO: Add description.
	generator (`None`, *optional*):
	TODO: Add description.
	num_images_per_prompt (`None`, *optional*, defaults to 1):
	TODO: Add description.
	image_latents (`list`, *optional*):
	TODO: Add description.
	latents (`Tensor | NoneType`):
	TODO: Add description.
	num_inference_steps (`None`):
	TODO: Add description.
	timesteps (`None`):
	TODO: Add description.
	sigmas (`None`, *optional*):
	TODO: Add description.
	guidance_scale (`None`, *optional*, defaults to 4.0):
	TODO: Add description.
	joint_attention_kwargs (`None`, *optional*):
	TODO: Add description.
	image_latent_ids (`Tensor`, *optional*):
	Position IDs for image latents. Shape: (B, img_seq_len, 4)
	output_type (`None`, *optional*, defaults to pil):
	TODO: Add description.

diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein.py

Lines 324 to 372 in 0f1abc4

	class Flux2KleinAutoBlocks(SequentialPipelineBlocks):
	"""
	Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein.
	Supported workflows:
	- `text2image`: requires `prompt`
	- `image_conditioned`: requires `image`, `prompt`
	Components:
	text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) image_processor (`Flux2ImageProcessor`)
	vae (`AutoencoderKLFlux2`) scheduler (`FlowMatchEulerDiscreteScheduler`) transformer
	(`Flux2Transformer2DModel`)
	Configs:
	is_distilled (default: True)
	Inputs:
	prompt (`None`, *optional*):
	TODO: Add description.
	max_sequence_length (`int`, *optional*, defaults to 512):
	TODO: Add description.
	text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
	TODO: Add description.
	image (`None`, *optional*):
	TODO: Add description.
	height (`None`, *optional*):
	TODO: Add description.
	width (`None`, *optional*):
	TODO: Add description.
	generator (`None`, *optional*):
	TODO: Add description.
	num_images_per_prompt (`None`, *optional*, defaults to 1):
	TODO: Add description.
	image_latents (`list`, *optional*):
	TODO: Add description.
	latents (`Tensor | NoneType`):
	TODO: Add description.
	num_inference_steps (`None`):
	TODO: Add description.
	timesteps (`None`):
	TODO: Add description.
	sigmas (`None`, *optional*):
	TODO: Add description.
	joint_attention_kwargs (`None`, *optional*):
	TODO: Add description.
	image_latent_ids (`Tensor`, *optional*):
	Position IDs for image latents. Shape: (B, img_seq_len, 4)
	output_type (`None`, *optional*, defaults to pil):
	TODO: Add description.

diffusers/src/diffusers/modular_pipelines/flux2/modular_blocks_flux2_klein_base.py

Lines 337 to 385 in 0f1abc4

	class Flux2KleinBaseAutoBlocks(SequentialPipelineBlocks):
	"""
	Auto blocks that perform the text-to-image and image-conditioned generation using Flux2-Klein (base model).
	Supported workflows:
	- `text2image`: requires `prompt`
	- `image_conditioned`: requires `image`, `prompt`
	Components:
	text_encoder (`Qwen3ForCausalLM`) tokenizer (`Qwen2TokenizerFast`) guider (`ClassifierFreeGuidance`)
	image_processor (`Flux2ImageProcessor`) vae (`AutoencoderKLFlux2`) scheduler
	(`FlowMatchEulerDiscreteScheduler`) transformer (`Flux2Transformer2DModel`)
	Configs:
	is_distilled (default: False)
	Inputs:
	prompt (`None`, *optional*):
	TODO: Add description.
	max_sequence_length (`int`, *optional*, defaults to 512):
	TODO: Add description.
	text_encoder_out_layers (`tuple`, *optional*, defaults to (9, 18, 27)):
	TODO: Add description.
	image (`None`, *optional*):
	TODO: Add description.
	height (`None`, *optional*):
	TODO: Add description.
	width (`None`, *optional*):
	TODO: Add description.
	generator (`None`, *optional*):
	TODO: Add description.
	num_images_per_prompt (`None`, *optional*, defaults to 1):
	TODO: Add description.
	latents (`Tensor | NoneType`):
	TODO: Add description.
	image_latents (`list`, *optional*):
	TODO: Add description.
	num_inference_steps (`None`):
	TODO: Add description.
	timesteps (`None`):
	TODO: Add description.
	sigmas (`None`, *optional*):
	TODO: Add description.
	joint_attention_kwargs (`None`, *optional*):
	TODO: Add description.
	image_latent_ids (`Tensor`, *optional*):
	Position IDs for image latents. Shape: (B, img_seq_len, 4)
	output_type (`None`, *optional*, defaults to pil):
	TODO: Add description.

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:

diffusers/tests/pipelines/flux2/test_pipeline_flux2.py

Lines 23 to 24 in 0f1abc4

	class Flux2PipelineFastTests(PipelineTesterMixin, unittest.TestCase):
	pipeline_class = Flux2Pipeline

diffusers/tests/pipelines/flux2/test_pipeline_flux2_klein.py

Lines 19 to 20 in 0f1abc4

	class Flux2KleinPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
	pipeline_class = Flux2KleinPipeline

diffusers/tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py

Lines 26 to 27 in 0f1abc4

	class Flux2KleinInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
	pipeline_class = Flux2KleinInpaintPipeline

diffusers/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py

Lines 19 to 20 in 0f1abc4

	class Flux2KleinKVPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
	pipeline_class = Flux2KleinKVPipeline

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py

Lines 38 to 71 in 0f1abc4

	class AutoencoderKLFlux2(
	ModelMixin, AutoencoderMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapterMixin
	):
	r"""
	A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
	This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
	for all models (such as downloading or saving).
	Parameters:
	in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
	out_channels (int, *optional*, defaults to 3): Number of channels in the output.
	down_block_types (`tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
	Tuple of downsample block types.
	up_block_types (`tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
	Tuple of upsample block types.
	block_out_channels (`tuple[int]`, *optional*, defaults to `(64,)`):
	Tuple of block output channels.
	act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
	latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
	sample_size (`int`, *optional*, defaults to `32`): Sample input size.
	force_upcast (`bool`, *optional*, default to `True`):
	If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
	can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
	can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
	mid_block_add_attention (`bool`, *optional*, default to `True`):
	If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
	mid_block will only have resnet blocks
	"""
	_supports_gradient_checkpointing = True
	_no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
	@register_to_config

diffusers/docs/source/en/api/pipelines/flux2.md

Lines 35 to 49 in 0f1abc4

	## Flux2Pipeline
	[[autodoc]] Flux2Pipeline
	- all
	- __call__
	## Flux2KleinPipeline
	[[autodoc]] Flux2KleinPipeline
	- all
	- __call__
	## Flux2KleinKVPipeline
	[[autodoc]] Flux2KleinKVPipeline

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:

diffusers/src/diffusers/training_utils.py

Lines 16 to 20 in 0f1abc4

	if getattr(torch, "distributed", None) is not None:
	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

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:

diffusers/src/diffusers/models/attention_dispatch.py

Lines 26 to 31 in 0f1abc4

	import torch.distributed as dist
	import torch.nn.functional as F
	if torch.distributed.is_available():
	import torch.distributed._functional_collectives as funcol

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