prx model/pipeline review

# `prx` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Files/categories reviewed: public exports/lazy imports, pipeline runtime, transformer attention/config behavior, docs, converter, fast tests, slow-test coverage, and duplicate GitHub Issues/PRs.

Duplicate search status: searched `prx`, `PRXPipeline`, `T5GemmaEncoder`, `callback_on_step_end`, `num_images_per_prompt`, `fuse_qkv_projections`, `latents dtype`, `attention_mask`, and `slow tests` across `huggingface/diffusers` Issues and PRs. Related existing items: closed issue https://github.com/huggingface/diffusers/issues/13142 and merged PR https://github.com/huggingface/diffusers/pull/13143 cover a composite `T5GemmaConfig` loading bug, but not the encoder-only config failure below. Open PR https://github.com/huggingface/diffusers/pull/13347 is related to PRX transformer test refactoring, not the pipeline gaps below.

## Issue 1: Encoder-only `T5GemmaEncoder` checkpoints still fail to load

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/__init__.py#L34-L43

Problem:
The wrapper only injects `config.encoder` when `T5GemmaConfig.from_pretrained(...)` returns a composite config with an `encoder` attribute. Encoder-only saves use `T5GemmaModuleConfig`, which is exactly what `T5GemmaEncoder.save_pretrained(...)` writes. In that case the wrapper falls through and `from_pretrained` instantiates with the wrong/default config.

Impact:
Converted or locally saved PRX pipelines can fail to reload their text encoder. This is related to #13142/#13143, but that fix is incomplete for encoder-only saved configs.

Reproduction:
```python
import tempfile
from diffusers.pipelines.prx import T5GemmaEncoder
from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig
from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder as RawT5GemmaEncoder

params = dict(vocab_size=16, hidden_size=8, intermediate_size=16, num_hidden_layers=1,
num_attention_heads=2, num_key_value_heads=1, head_dim=4,
max_position_embeddings=64, layer_types=["full_attention"])
config = T5GemmaConfig(encoder=T5GemmaModuleConfig(**params), is_encoder_decoder=False, **params)
model = RawT5GemmaEncoder(config.encoder)

with tempfile.TemporaryDirectory() as tmp:
model.save_pretrained(tmp)
T5GemmaEncoder.from_pretrained(tmp)
```

Relevant precedent:
`Wan`/`QwenImage` avoid custom import monkeypatching for their main components; where wrappers are unavoidable, the loader should handle both composite and component configs.

Suggested fix:
```python
from transformers import AutoConfig
from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig

config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
if isinstance(config, T5GemmaConfig) and hasattr(config, "encoder"):
kwargs["config"] = config.encoder
elif isinstance(config, T5GemmaModuleConfig):
kwargs["config"] = config
```

## Issue 2: CFG timestep shape breaks batched prompts and `num_images_per_prompt > 1`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L741-L749

Problem:
The CFG branch always builds `t_cont` with length `2`, but `latents_in` has length `2 * batch_size * num_images_per_prompt`.

Impact:
Any CFG run with more than one generated image in the effective batch crashes inside modulation broadcasting.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, PRXPipeline, PRXTransformer2DModel

transformer = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
pipe = PRXPipeline(transformer, FlowMatchEulerDiscreteScheduler(), None, None, None, 32)
pipe.set_progress_bar_config(disable=True)

pipe(prompt_embeds=torch.randn(2, 4, 8),
negative_prompt_embeds=torch.randn(2, 4, 8),
guidance_scale=2.0, height=32, width=32,
num_inference_steps=1, output_type="latent", use_resolution_binning=False)
```

Relevant precedent:
Flux/Qwen/Wan pipelines expand timestep tensors to the effective model input batch.

Suggested fix:
```python
latents_in = torch.cat([latents, latents], dim=0) if self.do_classifier_free_guidance else latents
t_cont = (t.float() / self.scheduler.config.num_train_timesteps).to(device)
t_cont = t_cont.reshape(1).expand(latents_in.shape[0])
```

## Issue 3: Step callbacks cannot modify tensors

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L768-L772

Problem:
`callback_on_step_end` is called, but its returned `callback_kwargs` are ignored.

Impact:
Standard diffusers callback behavior is broken. Users cannot edit latents/prompt embeddings during denoising, and the common callback mutation test is effectively bypassed by PRX’s custom test.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, PRXPipeline, PRXTransformer2DModel

transformer = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
pipe = PRXPipeline(transformer, FlowMatchEulerDiscreteScheduler(), None, None, None, 32)
pipe.set_progress_bar_config(disable=True)

def zero_latents(pipe, i, t, kwargs):
kwargs["latents"] = torch.zeros_like(kwargs["latents"])
return kwargs

out = pipe(prompt_embeds=torch.randn(1, 4, 8), guidance_scale=1.0,
height=32, width=32, num_inference_steps=1, output_type="latent",
use_resolution_binning=False, callback_on_step_end=zero_latents,
callback_on_step_end_tensor_inputs=["latents"])[0]
assert out.abs().sum() == 0
```

Relevant precedent:
`PipelineTesterMixin.test_callback_inputs` expects returned callback tensors to be applied.

Suggested fix:
```python
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
```

## Issue 4: User-provided latents are not cast to the pipeline dtype

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L346-L368

Problem:
Generated latents use the requested dtype, but supplied `latents` only move device and keep their original dtype.

Impact:
A bf16/fp16 pipeline can fail with a matmul dtype mismatch when users pass default fp32 latents.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, PRXPipeline, PRXTransformer2DModel

transformer = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
pipe = PRXPipeline(transformer, FlowMatchEulerDiscreteScheduler(), None, None, None, 32).to(dtype=torch.bfloat16)
pipe.set_progress_bar_config(disable=True)

pipe(prompt_embeds=torch.randn(1, 4, 8, dtype=torch.bfloat16),
latents=torch.randn(1, 4, 32, 32, dtype=torch.float32),
guidance_scale=1.0, height=32, width=32, num_inference_steps=1,
output_type="latent", use_resolution_binning=False)
```

Relevant precedent:
Most text-to-image pipelines cast supplied latents with `latents.to(device=device, dtype=dtype)`.

Suggested fix:
```python
else:
latents = latents.to(device=device, dtype=dtype)
```

## Issue 5: Precomputed attention masks are duplicated in the wrong order

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L394-L413

Problem:
`prompt_embeds` are repeated per prompt as `[p0, p0, p1, p1]`, but masks use `repeat(num_images_per_prompt, 1)`, producing `[m0, m1, m0, m1]`.

Impact:
For batched precomputed embeddings with `num_images_per_prompt > 1`, image copies receive the wrong text padding mask.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, PRXPipeline, PRXTransformer2DModel

transformer = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
pipe = PRXPipeline(transformer, FlowMatchEulerDiscreteScheduler(), None, None, None, 32)

mask = torch.tensor([[1, 0], [0, 1]], dtype=torch.bool)
embeds = torch.arange(2 * 2 * 8, dtype=torch.float32).reshape(2, 2, 8)
expanded, expanded_mask, *_ = pipe.encode_prompt(
None, device=torch.device("cpu"), do_classifier_free_guidance=False,
num_images_per_prompt=2, prompt_embeds=embeds, prompt_attention_mask=mask
)
print(expanded[:, 0, 0].tolist(), expanded_mask.tolist())
```

Relevant precedent:
Prompt/mask duplication should use the same batch ordering as embeddings.

Suggested fix:
```python
prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0)
negative_prompt_attention_mask = negative_prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0)
```

## Issue 6: `encode_prompt` forces no-grad

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L455-L460

Problem:
`_encode_prompt_standard` wraps the text encoder in `torch.no_grad()`. The pipeline `__call__` is already decorated with `@torch.no_grad()`.

Impact:
Calling `encode_prompt` directly with gradients enabled for embedding optimization or training cannot propagate gradients through the text encoder.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, PRXPipeline, PRXTransformer2DModel

class Tok:
model_max_length = 4
def __call__(self, texts, **kwargs):
return {"input_ids": torch.ones(len(texts), 4, dtype=torch.long),
"attention_mask": torch.ones(len(texts), 4, dtype=torch.long)}

class Enc(torch.nn.Module):
def __init__(self): super().__init__(); self.emb = torch.nn.Embedding(2, 8)
def forward(self, input_ids, **kwargs): return {"last_hidden_state": self.emb(input_ids)}

transformer = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
pipe = PRXPipeline(transformer, FlowMatchEulerDiscreteScheduler(), Enc(), Tok(), None, 32)

with torch.enable_grad():
prompt_embeds, *_ = pipe.encode_prompt("hello", device=torch.device("cpu"), do_classifier_free_guidance=False)
assert prompt_embeds.requires_grad
```

Relevant precedent:
Flux, QwenImage, StableAudio, and other pipelines rely on `__call__` for inference no-grad and keep prompt helpers grad-capable.

Suggested fix:
```python
embeddings = self.text_encoder(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
)["last_hidden_state"]
```

## Issue 7: Prompt cleaning hard-requires optional `ftfy`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L40-L41
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/prx/pipeline_prx.py#L199-L201

Problem:
`ftfy` is imported only when available, but `clean_text` calls `ftfy.fix_text` unconditionally.

Impact:
PRX prompt encoding fails in environments without the optional `ftfy` dependency.

Reproduction:
```python
from diffusers.pipelines.prx.pipeline_prx import TextPreprocessor
import diffusers.pipelines.prx.pipeline_prx as pipeline_prx

old_ftfy = getattr(pipeline_prx, "ftfy", None)
if hasattr(pipeline_prx, "ftfy"):
delattr(pipeline_prx, "ftfy")
try:
TextPreprocessor().clean_text("A prompt")
finally:
if old_ftfy is not None:
pipeline_prx.ftfy = old_ftfy
```

Relevant precedent:
`pipeline_wan.py` and `pipeline_kandinsky.py` guard `ftfy.fix_text` with `is_ftfy_available()`.

Suggested fix:
```python
if is_ftfy_available():
text = ftfy.fix_text(text)
text = html.unescape(html.unescape(text))
```

## Issue 8: Transformer exposes unsupported QKV fusion and lacks device-map split metadata

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_prx.py#L192-L199
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_prx.py#L590-L654

Problem:
`PRXAttention` inherits `_supports_qkv_fusion = True` from `AttentionModuleMixin`, but it uses `img_qkv_proj`/`txt_kv_proj` rather than `to_q`/`to_k`/`to_v`. The model also does not declare `_no_split_modules` or layerwise casting skip patterns.

Impact:
`fuse_qkv_projections()` raises `AttributeError`, and common device-map/offload tests skip PRX because split metadata is absent.

Reproduction:
```python
from diffusers import PRXTransformer2DModel

model = PRXTransformer2DModel(patch_size=1, in_channels=4, context_in_dim=8,
hidden_size=8, mlp_ratio=2.0, num_heads=2,
depth=1, axes_dim=[2, 2])
model.fuse_qkv_projections()
```

Relevant precedent:
Flux/QwenImage/Wan define `_no_split_modules` and `_skip_layerwise_casting_patterns`; Flux2 disables unsupported fusion for attention modules that are already fused.

Suggested fix:
```python
class PRXAttention(nn.Module, AttentionModuleMixin):
_supports_qkv_fusion = False
...

class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
_no_split_modules = ["PRXBlock"]
_skip_layerwise_casting_patterns = ["pe_embedder", "norm"]
```

## Issue 9: Slow tests are missing and key fast coverage is skipped

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/prx/test_pipeline_prx.py#L260-L282

Problem:
There are no PRX slow tests under `tests/`, and the PRX fast pipeline tests skip save/load, DDUF, variants, accelerator device-map, optional-component save/load, and dtype-dict loading.

Impact:
The loading and serialization regressions above are not caught by CI, and there is no published-checkpoint slow test for the supported Photoroom models.

Reproduction:
```python
from pathlib import Path

prx_tests = list(Path("tests").rglob("*prx*.py"))
slow_marked = [p for p in prx_tests if "@slow" in p.read_text(encoding="utf-8")]
assert slow_marked, f"No PRX slow tests found in: {[str(p) for p in prx_tests]}"
```

Relevant precedent:
Most mature pipelines keep fast component tests plus at least one slow `from_pretrained` smoke test for a public checkpoint.

Suggested fix:
Add a slow test that loads a small/public PRX checkpoint or a dedicated `hf-internal-testing` tiny PRX pipeline, runs a deterministic 1-2 step inference, and asserts shape/numerical slices. Re-enable serialization and dtype-dict tests once the T5Gemma loader is fixed.

Start with src/diffusers/pipelines/prx/__init__.py and pipeline_prx.py at the affected locations, then run the reproductions for checkpoint loading, batching, callbacks, dtype, masks, gradients, and optional ftfy. Compare callback behavior with PipelineTesterMixin.test_callback_inputs and related pipeline precedents named in the issue. Done means each reported PRX failure is fixed and covered by focused tests without regressing existing fast or slow coverage.

Vom Indexierungsmodell aus dem Issue-Text verfasst.