aura_flow model/pipeline review

[hlky]

aura_flow model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search performed with gh against huggingface/diffusers for AuraFlow, AuraFlowTransformer2DModel, AuraFlowPipeline, the specific failure modes below, and missing slow coverage.

Issue 1: Positional embedding indices can go out of bounds

Affected code:

diffusers/src/diffusers/models/transformers/auraflow_transformer_2d.py

Lines 72 to 108 in 0f1abc4

	def pe_selection_index_based_on_dim(self, h, w):
	# select subset of positional embedding based on H, W, where H, W is size of latent
	# PE will be viewed as 2d-grid, and H/p x W/p of the PE will be selected
	# because original input are in flattened format, we have to flatten this 2d grid as well.
	h_p, w_p = h // self.patch_size, w // self.patch_size
	h_max, w_max = int(self.pos_embed_max_size**0.5), int(self.pos_embed_max_size**0.5)
	# Calculate the top-left corner indices for the centered patch grid
	starth = h_max // 2 - h_p // 2
	startw = w_max // 2 - w_p // 2
	# Generate the row and column indices for the desired patch grid
	rows = torch.arange(starth, starth + h_p, device=self.pos_embed.device)
	cols = torch.arange(startw, startw + w_p, device=self.pos_embed.device)
	# Create a 2D grid of indices
	row_indices, col_indices = torch.meshgrid(rows, cols, indexing="ij")
	# Convert the 2D grid indices to flattened 1D indices
	selected_indices = (row_indices * w_max + col_indices).flatten()
	return selected_indices
	def forward(self, latent) -> torch.Tensor:
	batch_size, num_channels, height, width = latent.size()
	latent = latent.view(
	batch_size,
	num_channels,
	height // self.patch_size,
	self.patch_size,
	width // self.patch_size,
	self.patch_size,
	)
	latent = latent.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
	latent = self.proj(latent)
	pe_index = self.pe_selection_index_based_on_dim(height, width)
	return latent + self.pos_embed[:, pe_index]

Problem:
This is already tracked by #12656 and open PR #13110. When latent spatial dimensions exceed the learned PE grid, pe_selection_index_based_on_dim() creates negative or too-large indices and fails with an indexing error, potentially as a CUDA device assert.

Impact:
Users requesting larger resolutions or training/fine-tuning at larger latent sizes get a low-level crash instead of a clear validation error.

Reproduction:

import torch
from diffusers import AuraFlowTransformer2DModel

model = AuraFlowTransformer2DModel(sample_size=4, patch_size=2, in_channels=4, out_channels=4, num_mmdit_layers=0, num_single_dit_layers=0, attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8, caption_projection_dim=4, pos_embed_max_size=4)
model(hidden_states=torch.randn(1, 4, 8, 8), encoder_hidden_states=torch.randn(1, 2, 8), timestep=torch.tensor([1.0]))

Relevant precedent:
PR #13110 adds the right kind of bounds check.

Suggested fix:

if h_p > h_max or w_p > w_max:
 raise ValueError(
 f"Input patch grid ({h_p}, {w_p}) exceeds AuraFlow positional embedding grid ({h_max}, {w_max})."
 )

Issue 2: out_channels=None is serialized but crashes in forward

Affected code:

diffusers/src/diffusers/models/transformers/auraflow_transformer_2d.py

Lines 319 to 320 in 0f1abc4

	default_out_channels = in_channels
	self.out_channels = out_channels if out_channels is not None else default_out_channels

diffusers/src/diffusers/models/transformers/auraflow_transformer_2d.py

Lines 461 to 468 in 0f1abc4

	# unpatchify
	patch_size = self.config.patch_size
	out_channels = self.config.out_channels
	height = height // patch_size
	width = width // patch_size
	hidden_states = hidden_states.reshape(
	shape=(hidden_states.shape[0], height, width, patch_size, patch_size, out_channels)

Problem:
__init__ handles out_channels=None by setting self.out_channels, but forward() uses self.config.out_channels, which remains None.

Impact:
A valid constructor path and any config that stores out_channels: null crashes during unpatchify.

Reproduction:

import torch
from diffusers import AuraFlowTransformer2DModel

model = AuraFlowTransformer2DModel(sample_size=4, patch_size=2, in_channels=4, out_channels=None, num_mmdit_layers=0, num_single_dit_layers=0, attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8, caption_projection_dim=4, pos_embed_max_size=4)
model(hidden_states=torch.randn(1, 4, 4, 4), encoder_hidden_states=torch.randn(1, 2, 8), timestep=torch.tensor([1.0]))

Relevant precedent:
SD3Transformer2DModel unpatchifies with self.out_channels.

Suggested fix:

out_channels = self.out_channels

Issue 3: Provided prompt masks are required but not applied

Affected code:

diffusers/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

Lines 312 to 332 in 0f1abc4

	text_inputs = {k: v.to(device) for k, v in text_inputs.items()}
	prompt_embeds = self.text_encoder(**text_inputs)[0]
	prompt_attention_mask = text_inputs["attention_mask"].unsqueeze(-1).expand(prompt_embeds.shape)
	prompt_embeds = prompt_embeds * prompt_attention_mask
	if self.text_encoder is not None:
	dtype = self.text_encoder.dtype
	elif self.transformer is not None:
	dtype = self.transformer.dtype
	else:
	dtype = None
	prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
	bs_embed, seq_len, _ = prompt_embeds.shape
	# duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
	prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
	prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
	prompt_attention_mask = prompt_attention_mask.reshape(bs_embed, -1)
	prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)

diffusers/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

Lines 622 to 628 in 0f1abc4

	noise_pred = self.transformer(
	latent_model_input,
	encoder_hidden_states=prompt_embeds,
	timestep=timestep,
	return_dict=False,
	attention_kwargs=self.attention_kwargs,
	)[0]

Problem:
For tokenizer-generated embeddings, the pipeline zeroes padded embeddings. For user-supplied prompt_embeds, prompt_attention_mask is required but only reshaped/repeated, then never passed to the transformer. This is related to closed issue #8886; the current precomputed-embedding path still reproduces the same mask-semantics gap.

Impact:
Two prompt_embeds that differ only in masked positions can condition the model differently, so precomputed embeddings are not equivalent to the tokenizer path.

Reproduction:

import torch
from diffusers import AuraFlowPipeline

pipe = AuraFlowPipeline(tokenizer=None, text_encoder=None, vae=None, transformer=None, scheduler=None)
embeds = torch.ones(1, 4, 8)
embeds[:, 2:, :] = 100.0
mask = torch.tensor([[1, 1, 0, 0]])

prompt_embeds, returned_mask, _, _ = pipe.encode_prompt(
 prompt=None,
 prompt_embeds=embeds,
 prompt_attention_mask=mask,
 do_classifier_free_guidance=False,
)
print(prompt_embeds[0, 2:].abs().max().item()) # 100.0

Relevant precedent:
QwenImage passes a real encoder attention mask through to the transformer; PixArt passes encoder_attention_mask into the denoiser.

Suggested fix:

prompt_attention_mask = prompt_attention_mask.to(device=device)
prompt_embeds = prompt_embeds * prompt_attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype)

Apply the same handling for negative_prompt_embeds, and keep masks 2D if they are later passed into attention.

Issue 4: fp16 VAE upcasting mutates the pipeline permanently

Affected code:

diffusers/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py

Lines 657 to 662 in 0f1abc4

	# make sure the VAE is in float32 mode, as it overflows in float16
	needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
	if needs_upcasting:
	self.upcast_vae()
	latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
	image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]

Problem:
The pipeline calls deprecated self.upcast_vae() internally and never casts the VAE back to its original dtype.

Impact:
A fp16 pipeline silently keeps the VAE in fp32 after the first decoded call, increasing memory use and emitting an internal deprecation warning.

Reproduction:

import torch
from diffusers import AutoencoderKL, AuraFlowPipeline, AuraFlowTransformer2DModel, FlowMatchEulerDiscreteScheduler

vae = AutoencoderKL(block_out_channels=[4], in_channels=3, out_channels=3, down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"], latent_channels=4, sample_size=8, norm_num_groups=1)
vae.config.force_upcast = True
vae.to(dtype=torch.float16)

transformer = AuraFlowTransformer2DModel(sample_size=8, patch_size=1, in_channels=4, out_channels=4, num_mmdit_layers=0, num_single_dit_layers=0, attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8, caption_projection_dim=4, pos_embed_max_size=64)
pipe = AuraFlowPipeline(None, None, vae, transformer, FlowMatchEulerDiscreteScheduler())
pipe.set_progress_bar_config(disable=True)

print(pipe.vae.dtype)
pipe(prompt=None, prompt_embeds=torch.zeros(1, 4, 8), prompt_attention_mask=torch.ones(1, 4), latents=torch.zeros(1, 4, 8, 8), height=8, width=8, guidance_scale=1.0, num_inference_steps=1, output_type="np")
print(pipe.vae.dtype) # torch.float32

Relevant precedent:
Stable Diffusion XL casts the VAE back after decoding.

Suggested fix:

vae_dtype = self.vae.dtype
if needs_upcasting:
 self.vae.to(dtype=torch.float32)
 latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]

if needs_upcasting:
 self.vae.to(dtype=vae_dtype)

Issue 5: AuraFlow attention processors ignore attention backend dispatch

Affected code:

diffusers/src/diffusers/models/attention_processor.py

Lines 2087 to 2156 in 0f1abc4

	class AuraFlowAttnProcessor2_0:
	"""Attention processor used typically in processing Aura Flow."""
	def __init__(self):
	if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"):
	raise ImportError(
	"AuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. "
	)
	def __call__(
	self,
	attn: Attention,
	hidden_states: torch.FloatTensor,
	encoder_hidden_states: torch.FloatTensor = None,
	*args,
	**kwargs,
	) -> torch.FloatTensor:
	batch_size = hidden_states.shape[0]
	# `sample` projections.
	query = attn.to_q(hidden_states)
	key = attn.to_k(hidden_states)
	value = attn.to_v(hidden_states)
	# `context` projections.
	if encoder_hidden_states is not None:
	encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
	encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
	encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
	# Reshape.
	inner_dim = key.shape[-1]
	head_dim = inner_dim // attn.heads
	query = query.view(batch_size, -1, attn.heads, head_dim)
	key = key.view(batch_size, -1, attn.heads, head_dim)
	value = value.view(batch_size, -1, attn.heads, head_dim)
	# Apply QK norm.
	if attn.norm_q is not None:
	query = attn.norm_q(query)
	if attn.norm_k is not None:
	key = attn.norm_k(key)
	# Concatenate the projections.
	if encoder_hidden_states is not None:
	encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
	batch_size, -1, attn.heads, head_dim
	)
	encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)
	encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
	batch_size, -1, attn.heads, head_dim
	)
	if attn.norm_added_q is not None:
	encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
	if attn.norm_added_k is not None:
	encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
	query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)
	key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
	value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
	query = query.transpose(1, 2)
	key = key.transpose(1, 2)
	value = value.transpose(1, 2)
	# Attention.
	hidden_states = F.scaled_dot_product_attention(
	query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False
	)

diffusers/src/diffusers/models/attention_processor.py

Lines 2180 to 2253 in 0f1abc4

	class FusedAuraFlowAttnProcessor2_0:
	"""Attention processor used typically in processing Aura Flow with fused projections."""
	def __init__(self):
	if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"):
	raise ImportError(
	"FusedAuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. "
	)
	def __call__(
	self,
	attn: Attention,
	hidden_states: torch.FloatTensor,
	encoder_hidden_states: torch.FloatTensor = None,
	*args,
	**kwargs,
	) -> torch.FloatTensor:
	batch_size = hidden_states.shape[0]
	# `sample` projections.
	qkv = attn.to_qkv(hidden_states)
	split_size = qkv.shape[-1] // 3
	query, key, value = torch.split(qkv, split_size, dim=-1)
	# `context` projections.
	if encoder_hidden_states is not None:
	encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
	split_size = encoder_qkv.shape[-1] // 3
	(
	encoder_hidden_states_query_proj,
	encoder_hidden_states_key_proj,
	encoder_hidden_states_value_proj,
	) = torch.split(encoder_qkv, split_size, dim=-1)
	# Reshape.
	inner_dim = key.shape[-1]
	head_dim = inner_dim // attn.heads
	query = query.view(batch_size, -1, attn.heads, head_dim)
	key = key.view(batch_size, -1, attn.heads, head_dim)
	value = value.view(batch_size, -1, attn.heads, head_dim)
	# Apply QK norm.
	if attn.norm_q is not None:
	query = attn.norm_q(query)
	if attn.norm_k is not None:
	key = attn.norm_k(key)
	# Concatenate the projections.
	if encoder_hidden_states is not None:
	encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
	batch_size, -1, attn.heads, head_dim
	)
	encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)
	encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
	batch_size, -1, attn.heads, head_dim
	)
	if attn.norm_added_q is not None:
	encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
	if attn.norm_added_k is not None:
	encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
	query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)
	key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
	value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
	query = query.transpose(1, 2)
	key = key.transpose(1, 2)
	value = value.transpose(1, 2)
	# Attention.
	hidden_states = F.scaled_dot_product_attention(
	query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False
	)

Problem:
AuraFlowAttnProcessor2_0 and FusedAuraFlowAttnProcessor2_0 call F.scaled_dot_product_attention directly and do not expose _attention_backend, so model.set_attention_backend(...) is effectively a no-op for AuraFlow processors. I did not find an existing duplicate for this; PR #13533 fixed a different AuraFlow attention-processor bug.

Impact:
AuraFlow misses the repo's current backend dispatch behavior for flash/flex/sage/native variants and context-parallel plumbing.

Reproduction:

from diffusers import AuraFlowTransformer2DModel

model = AuraFlowTransformer2DModel(sample_size=2, patch_size=1, in_channels=4, num_mmdit_layers=1, num_single_dit_layers=1, attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8, caption_projection_dim=4, pos_embed_max_size=4)
model.set_attention_backend("native")
print(sum(hasattr(p, "_attention_backend") for p in model.attn_processors.values())) # 0

Relevant precedent:
FluxAttnProcessor in transformer_flux.py defines _attention_backend / _parallel_config and calls dispatch_attention_fn.

Suggested fix:
Move/refactor AuraFlow processors to the transformer file or update them in place to follow the current processor contract:

hidden_states = dispatch_attention_fn(
 query,
 key,
 value,
 attn_mask=attention_mask,
 backend=self._attention_backend,
 parallel_config=self._parallel_config,
)

Issue 6: No dedicated slow AuraFlow pipeline test

Affected code:

diffusers/tests/pipelines/aura_flow/test_pipeline_aura_flow.py

Lines 16 to 137 in 0f1abc4

	class AuraFlowPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
	pipeline_class = AuraFlowPipeline
	params = frozenset(
	[
	"prompt",
	"height",
	"width",
	"guidance_scale",
	"negative_prompt",
	"prompt_embeds",
	"negative_prompt_embeds",
	]
	)
	batch_params = frozenset(["prompt", "negative_prompt"])
	test_layerwise_casting = True
	test_group_offloading = True
	def get_dummy_components(self):
	torch.manual_seed(0)
	transformer = AuraFlowTransformer2DModel(
	sample_size=32,
	patch_size=2,
	in_channels=4,
	num_mmdit_layers=1,
	num_single_dit_layers=1,
	attention_head_dim=8,
	num_attention_heads=4,
	caption_projection_dim=32,
	joint_attention_dim=32,
	out_channels=4,
	pos_embed_max_size=256,
	)
	text_encoder = UMT5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-umt5")
	tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
	torch.manual_seed(0)
	vae = AutoencoderKL(
	block_out_channels=[32, 64],
	in_channels=3,
	out_channels=3,
	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],
	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
	latent_channels=4,
	sample_size=32,
	)
	scheduler = FlowMatchEulerDiscreteScheduler()
	return {
	"scheduler": scheduler,
	"text_encoder": text_encoder,
	"tokenizer": tokenizer,
	"transformer": transformer,
	"vae": vae,
	}
	def get_dummy_inputs(self, device, seed=0):
	if str(device).startswith("mps"):
	generator = torch.manual_seed(seed)
	else:
	generator = torch.Generator(device="cpu").manual_seed(seed)
	inputs = {
	"prompt": "A painting of a squirrel eating a burger",
	"generator": generator,
	"num_inference_steps": 2,
	"guidance_scale": 5.0,
	"output_type": "np",
	"height": None,
	"width": None,
	}
	return inputs
	def test_attention_slicing_forward_pass(self):
	# Attention slicing needs to implemented differently for this because how single DiT and MMDiT
	# blocks interfere with each other.
	return
	def test_fused_qkv_projections(self):
	device = "cpu" # ensure determinism for the device-dependent torch.Generator
	components = self.get_dummy_components()
	pipe = self.pipeline_class(**components)
	pipe = pipe.to(device)
	pipe.set_progress_bar_config(disable=None)
	inputs = self.get_dummy_inputs(device)
	image = pipe(**inputs).images
	original_image_slice = image[0, -3:, -3:, -1]
	# TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
	# to the pipeline level.
	pipe.transformer.fuse_qkv_projections()
	assert check_qkv_fusion_processors_exist(pipe.transformer), (
	"Something wrong with the fused attention processors. Expected all the attention processors to be fused."
	)
	assert check_qkv_fusion_matches_attn_procs_length(
	pipe.transformer, pipe.transformer.original_attn_processors
	), "Something wrong with the attention processors concerning the fused QKV projections."
	inputs = self.get_dummy_inputs(device)
	image = pipe(**inputs).images
	image_slice_fused = image[0, -3:, -3:, -1]
	pipe.transformer.unfuse_qkv_projections()
	inputs = self.get_dummy_inputs(device)
	image = pipe(**inputs).images
	image_slice_disabled = image[0, -3:, -3:, -1]
	assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
	"Fusion of QKV projections shouldn't affect the outputs."
	)
	assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), (
	"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
	)
	assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
	"Original outputs should match when fused QKV projections are disabled."
	)
	@unittest.skip("xformers attention processor does not exist for AuraFlow")
	def test_xformers_attention_forwardGenerator_pass(self):
	pass

Problem:
AuraFlow has fast model/pipeline tests, LoRA tests, and a GGUF nightly path, but no dedicated slow/full-checkpoint AuraFlowPipeline integration test under tests/pipelines/aura_flow.

Impact:
Full checkpoint behavior for the standard pipeline can regress without a model-slice assertion. This is especially relevant for prompt embedding/mask behavior, VAE decode dtype handling, scheduler defaults, and resolution behavior.

Reproduction:

from pathlib import Path

files = sorted(Path("tests/pipelines/aura_flow").glob("test_*.py"))
slow_hits = [str(p) for p in files if "@slow" in p.read_text() or "SlowTests" in p.read_text()]
print(slow_hits) # []

Relevant precedent:
tests/pipelines/flux/test_pipeline_flux.py and tests/pipelines/pixart_alpha/test_pixart.py include slow integration coverage with expected output slices.

Suggested fix:
Add an AuraFlowPipelineSlowTests class loading fal/AuraFlow-v0.3 or a maintained test-slice fixture, running 1-2 denoising steps with a fixed seed, and asserting a stable output slice.

Validation note: I attempted the AuraFlow pytest files with .venv, but collection fails in this environment because the installed Torch build lacks torch._C._distributed_c10d, imported via diffusers.training_utils. Standalone reproductions above were run with .venv.