Skip to content

prx model/pipeline review #13608

Open
Open
prx model/pipeline review#13608

Description

@hlky
hlky
opened on Apr 29, 2026

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions