aura_flow model/pipeline review

# `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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/auraflow_transformer_2d.py#L72-L108

Problem:
This is already tracked by https://github.com/huggingface/diffusers/issues/12656 and open PR https://github.com/huggingface/diffusers/pull/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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/auraflow_transformer_2d.py#L319-L320
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/auraflow_transformer_2d.py#L461-L468

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:
```python
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:
```python
out_channels = self.out_channels
```

## Issue 3: Provided prompt masks are required but not applied

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py#L312-L332
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py#L622-L628

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 https://github.com/huggingface/diffusers/issues/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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py#L657-L662

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L2087-L2156
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L2180-L2253

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/aura_flow/test_pipeline_aura_flow.py#L16-L137

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:
```python
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`.

Start by separating the six findings and reading the affected files: src/diffusers/models/transformers/auraflow_transformer_2d.py, src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py, src/diffusers/models/attention_processor.py, and tests/pipelines/aura_flow/test_pipeline_aura_flow.py. Run the supplied reproductions and existing AuraFlow tests; done means each confirmed defect has coverage and the slow pipeline integration test is present, while Issue 1 should be checked against #12656 and PR #13110.

Written by the indexing model from the issue text.