prx model/pipeline review

[hlky]

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 #13142 and merged PR #13143 cover a composite T5GemmaConfig loading bug, but not the encoder-only config failure below. Open PR #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:

diffusers/src/diffusers/pipelines/prx/__init__.py

Lines 34 to 43 in 0f1abc4

	class T5GemmaEncoder(_T5GemmaEncoder):
	@classmethod
	def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
	if "config" not in kwargs:
	from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig
	config = T5GemmaConfig.from_pretrained(pretrained_model_name_or_path)
	if hasattr(config, "encoder"):
	kwargs["config"] = config.encoder
	return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)

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:

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:

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:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 741 to 749 in 0f1abc4

	# Duplicate latents if using classifier-free guidance
	if self.do_classifier_free_guidance:
	latents_in = torch.cat([latents, latents], dim=0)
	# Normalize timestep for the transformer
	t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).repeat(2).to(device)
	else:
	latents_in = latents
	# Normalize timestep for the transformer
	t_cont = (t.float() / self.scheduler.config.num_train_timesteps).view(1).to(device)

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:

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:

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:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 768 to 772 in 0f1abc4

	if callback_on_step_end is not None:
	callback_kwargs = {}
	for k in callback_on_step_end_tensor_inputs:
	callback_kwargs[k] = locals()[k]
	callback_on_step_end(self, i, t, callback_kwargs)

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:

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:

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:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 346 to 368 in 0f1abc4

	def prepare_latents(
	self,
	batch_size: int,
	num_channels_latents: int,
	height: int,
	width: int,
	dtype: torch.dtype,
	device: torch.device,
	generator: torch.Generator | None = None,
	latents: torch.Tensor | None = None,
	):
	"""Prepare initial latents for the diffusion process."""
	if latents is None:
	spatial_compression = self.vae_scale_factor
	latent_height, latent_width = (
	height // spatial_compression,
	width // spatial_compression,
	)
	shape = (batch_size, num_channels_latents, latent_height, latent_width)
	latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
	else:
	latents = latents.to(device)
	return latents

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:

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:

else:
 latents = latents.to(device=device, dtype=dtype)

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

Affected code:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 394 to 413 in 0f1abc4

	# Duplicate embeddings for each generation per prompt
	if num_images_per_prompt > 1:
	# Repeat prompt embeddings
	bs_embed, seq_len, _ = prompt_embeds.shape
	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)
	if prompt_attention_mask is not None:
	prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
	prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
	# Repeat negative embeddings if using CFG
	if do_classifier_free_guidance and negative_prompt_embeds is not None:
	bs_embed, seq_len, _ = negative_prompt_embeds.shape
	negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
	negative_prompt_embeds = negative_prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
	if negative_prompt_attention_mask is not None:
	negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
	negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)

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:

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:

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:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 455 to 460 in 0f1abc4

	with torch.no_grad():
	embeddings = self.text_encoder(
	input_ids=input_ids,
	attention_mask=attention_mask,
	output_hidden_states=True,
	)["last_hidden_state"]

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:

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:

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:

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 40 to 41 in 0f1abc4

	if is_ftfy_available():
	import ftfy

diffusers/src/diffusers/pipelines/prx/pipeline_prx.py

Lines 199 to 201 in 0f1abc4

	# Basic cleaning
	text = ftfy.fix_text(text)
	text = html.unescape(html.unescape(text))

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:

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:

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:

diffusers/src/diffusers/models/transformers/transformer_prx.py

Lines 192 to 199 in 0f1abc4

	class PRXAttention(nn.Module, AttentionModuleMixin):
	r"""
	PRX-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for
	PRX's architecture.
	"""
	_default_processor_cls = PRXAttnProcessor2_0
	_available_processors = [PRXAttnProcessor2_0]

diffusers/src/diffusers/models/transformers/transformer_prx.py

Lines 590 to 654 in 0f1abc4

	class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
	r"""
	Transformer-based 2D model for text to image generation.
	Args:
	in_channels (`int`, *optional*, defaults to 16):
	Number of input channels in the latent image.
	patch_size (`int`, *optional*, defaults to 2):
	Size of the square patches used to flatten the input image.
	context_in_dim (`int`, *optional*, defaults to 2304):
	Dimensionality of the text conditioning input.
	hidden_size (`int`, *optional*, defaults to 1792):
	Dimension of the hidden representation.
	mlp_ratio (`float`, *optional*, defaults to 3.5):
	Expansion ratio for the hidden dimension inside MLP blocks.
	num_heads (`int`, *optional*, defaults to 28):
	Number of attention heads.
	depth (`int`, *optional*, defaults to 16):
	Number of transformer blocks.
	axes_dim (`list[int]`, *optional*):
	list of dimensions for each positional embedding axis. Defaults to `[32, 32]`.
	theta (`int`, *optional*, defaults to 10000):
	Frequency scaling factor for rotary embeddings.
	time_factor (`float`, *optional*, defaults to 1000.0):
	Scaling factor applied in timestep embeddings.
	time_max_period (`int`, *optional*, defaults to 10000):
	Maximum frequency period for timestep embeddings.
	Attributes:
	pe_embedder (`EmbedND`):
	Multi-axis rotary embedding generator for positional encodings.
	img_in (`nn.Linear`):
	Projection layer for image patch tokens.
	time_in (`MLPEmbedder`):
	Embedding layer for timestep embeddings.
	txt_in (`nn.Linear`):
	Projection layer for text conditioning.
	blocks (`nn.ModuleList`):
	Stack of transformer blocks (`PRXBlock`).
	final_layer (`LastLayer`):
	Projection layer mapping hidden tokens back to patch outputs.
	Methods:
	attn_processors:
	Returns a dictionary of all attention processors in the model.
	set_attn_processor(processor):
	Replaces attention processors across all attention layers.
	process_inputs(image_latent, txt):
	Converts inputs into patch tokens, encodes text, and produces positional encodings.
	compute_timestep_embedding(timestep, dtype):
	Creates a timestep embedding of dimension 256, scaled and projected.
	forward_transformers(image_latent, cross_attn_conditioning, timestep, time_embedding, attention_mask,
	**block_kwargs):
	Runs the sequence of transformer blocks over image and text tokens.
	forward(image_latent, timestep, cross_attn_conditioning, micro_conditioning, cross_attn_mask=None,
	attention_kwargs=None, return_dict=True):
	Full forward pass from latent input to reconstructed output image.
	Returns:
	`Transformer2DModelOutput` if `return_dict=True` (default), otherwise a tuple containing:
	- `sample` (`torch.Tensor`): Reconstructed image of shape `(B, C, H, W)`.
	"""
	config_name = "config.json"
	_supports_gradient_checkpointing = True

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:

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:

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:

diffusers/tests/pipelines/prx/test_pipeline_prx.py

Lines 260 to 282 in 0f1abc4

	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_save_load_dduf(self):
	pass
	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_loading_with_variants(self):
	pass
	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_pipeline_with_accelerator_device_map(self):
	pass
	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_save_load_local(self):
	pass
	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_save_load_optional_components(self):
	pass
	@unittest.skip("Custom T5GemmaEncoder not compatible with transformers v5.")
	def test_torch_dtype_dict(self):
	pass

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:

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.