Add PRXPixelPipeline: pixel-space PRX text-to-image pipeline · huggingface/diffusers@71fcc14 · GitHub
Skip to content

Commit 71fcc14

Browse files
DavidBertclaude
andcommitted
Add PRXPixelPipeline: pixel-space PRX text-to-image pipeline
Adds a pixel-space variant of PRX that denoises raw RGB directly (no VAE), conditioned on a Qwen3-VL text encoder: - PRXTransformer2DModel: new optional config args `bottleneck_size` (two-layer img_in projection for large patch dims) and `resolution_embeds` (PRXResolutionEmbedder conditions the timestep modulation on the latent resolution) - PRXPipeline: support for subclass-tuned tokenizer max length, light text cleaning, x-prediction flow matching (x0 -> velocity conversion), and non-unit initial noise scale - PRXPixelPipeline: thin subclass wiring the above together (vae optional/None, vae_scale_factor=1, 1024px default) - conversion script support for the pixel checkpoint format - registration in __init__ files + dummy objects, docs autodoc entry, fast pipeline tests Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent d6c360d commit 71fcc14

10 files changed

Lines changed: 1393 additions & 196 deletions

File tree

‎docs/source/en/_toctree.yml‎

Lines changed: 2 additions & 0 deletions

‎scripts/convert_prx_to_diffusers.py‎

Lines changed: 256 additions & 177 deletions
Large diffs are not rendered by default.

‎src/diffusers/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,7 @@
671671
"PixArtSigmaPAGPipeline",
672672
"PixArtSigmaPipeline",
673673
"PRXPipeline",
674+
"PRXPixelPipeline",
674675
"QwenImageControlNetInpaintPipeline",
675676
"QwenImageControlNetPipeline",
676677
"QwenImageEditInpaintPipeline",
@@ -1491,6 +1492,7 @@
14911492
PixArtSigmaPAGPipeline,
14921493
PixArtSigmaPipeline,
14931494
PRXPipeline,
1495+
PRXPixelPipeline,
14941496
QwenImageControlNetInpaintPipeline,
14951497
QwenImageControlNetPipeline,
14961498
QwenImageEditInpaintPipeline,

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

Lines changed: 74 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from typing import Any
15+
from typing import Any, Optional
1616

1717
import torch
1818
from torch import nn
@@ -321,6 +321,48 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
321321
return self.out_layer(self.silu(self.in_layer(x)))
322322

323323

324+
class PRXResolutionEmbedder(nn.Module):
325+
r"""
326+
Embeds the spatial resolution `(height, width)` of the latent into a vector that is added to the timestep
327+
embedding, so the model can condition its modulation on the generation resolution.
328+
329+
A sinusoidal embedding of dimension 128 is built for the height and the width separately and concatenated into a
330+
256-dim vector, which is then projected to `hidden_size` by a 2-layer MLP. This matches the `"vec"` mode of the
331+
resolution-aware conditioning used during PRX-7B training.
332+
333+
Args:
334+
hidden_size (`int`):
335+
Dimension of the output embedding (must match the timestep embedding dimension).
336+
max_period (`int`, *optional*, defaults to 10000):
337+
Maximum frequency period for the sinusoidal resolution embedding.
338+
"""
339+
340+
def __init__(self, hidden_size: int, max_period: int = 10000):
341+
super().__init__()
342+
self.max_period = max_period
343+
self.mlp = MLPEmbedder(in_dim=256, hidden_dim=hidden_size)
344+
345+
def forward(self, height: torch.Tensor, width: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
346+
h_emb = get_timestep_embedding(
347+
timesteps=height,
348+
embedding_dim=128,
349+
max_period=self.max_period,
350+
scale=1.0,
351+
flip_sin_to_cos=True,
352+
downscale_freq_shift=0.0,
353+
)
354+
w_emb = get_timestep_embedding(
355+
timesteps=width,
356+
embedding_dim=128,
357+
max_period=self.max_period,
358+
scale=1.0,
359+
flip_sin_to_cos=True,
360+
downscale_freq_shift=0.0,
361+
)
362+
hw_emb = torch.cat([h_emb, w_emb], dim=-1).to(dtype)
363+
return self.mlp(hw_emb)
364+
365+
324366
class Modulation(nn.Module):
325367
r"""
326368
Modulation network that generates scale, shift, and gating parameters.
@@ -613,12 +655,19 @@ class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
613655
Scaling factor applied in timestep embeddings.
614656
time_max_period (`int`, *optional*, defaults to 10000):
615657
Maximum frequency period for timestep embeddings.
658+
bottleneck_size (`int`, *optional*):
659+
If set, the image patch projection (`img_in`) uses a two-layer bottleneck (`patch_dim -> bottleneck_size ->
660+
hidden_size`) instead of a single linear layer. Used by the pixel-space PRX-7B variant where the patch
661+
dimension is large.
662+
resolution_embeds (`bool`, *optional*, defaults to `False`):
663+
Whether to condition the timestep modulation on the latent resolution `(H, W)` via a
664+
`PRXResolutionEmbedder`. Used by the PRX-7B variant.
616665
617666
Attributes:
618667
pe_embedder (`EmbedND`):
619668
Multi-axis rotary embedding generator for positional encodings.
620-
img_in (`nn.Linear`):
621-
Projection layer for image patch tokens.
669+
img_in (`nn.Linear` or `nn.Sequential`):
670+
Projection layer for image patch tokens (a two-layer bottleneck when `bottleneck_size` is set).
622671
time_in (`MLPEmbedder`):
623672
Embedding layer for timestep embeddings.
624673
txt_in (`nn.Linear`):
@@ -666,6 +715,8 @@ def __init__(
666715
theta: int = 10000,
667716
time_factor: float = 1000.0,
668717
time_max_period: int = 10000,
718+
bottleneck_size: Optional[int] = None,
719+
resolution_embeds: bool = False,
669720
):
670721
super().__init__()
671722

@@ -691,10 +742,22 @@ def __init__(
691742
self.hidden_size = hidden_size
692743
self.num_heads = num_heads
693744
self.pe_embedder = PRXEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
694-
self.img_in = nn.Linear(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
745+
patch_dim = self.in_channels * self.patch_size**2
746+
if bottleneck_size is not None:
747+
# Two-layer bottleneck projection (used by pixel-space PRX where the patch dimension is large).
748+
self.img_in = nn.Sequential(
749+
nn.Linear(patch_dim, bottleneck_size, bias=True),
750+
nn.Linear(bottleneck_size, self.hidden_size, bias=True),
751+
)
752+
else:
753+
self.img_in = nn.Linear(patch_dim, self.hidden_size, bias=True)
695754
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
696755
self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
697756

757+
self.resolution_embedder = (
758+
PRXResolutionEmbedder(self.hidden_size, max_period=time_max_period) if resolution_embeds else None
759+
)
760+
698761
self.blocks = nn.ModuleList(
699762
[
700763
PRXBlock(
@@ -771,6 +834,13 @@ def forward(
771834
# Compute time embedding
772835
vec = self._compute_timestep_embedding(timestep, dtype=img.dtype)
773836

837+
# Add resolution conditioning (PRX-7B "vec" mode): embed the latent (H, W) and add it to the timestep vector
838+
# so every block's modulation is resolution-aware.
839+
if self.resolution_embedder is not None:
840+
height = torch.full((bs,), h, device=hidden_states.device, dtype=torch.float32)
841+
width = torch.full((bs,), w, device=hidden_states.device, dtype=torch.float32)
842+
vec = vec + self.resolution_embedder(height, width, dtype=vec.dtype)
843+
774844
# Apply transformer blocks
775845
for block in self.blocks:
776846
if torch.is_grad_enabled() and self.gradient_checkpointing:

‎src/diffusers/pipelines/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@
193193
"FluxKontextPipeline",
194194
"FluxKontextInpaintPipeline",
195195
]
196-
_import_structure["prx"] = ["PRXPipeline"]
196+
_import_structure["prx"] = ["PRXPipeline", "PRXPixelPipeline"]
197197
_import_structure["audioldm2"] = [
198198
"AudioLDM2Pipeline",
199199
"AudioLDM2ProjectionModel",
@@ -848,7 +848,7 @@
848848
StableDiffusionXLPAGPipeline,
849849
)
850850
from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline
851-
from .prx import PRXPipeline
851+
from .prx import PRXPipeline, PRXPixelPipeline
852852
from .qwenimage import (
853853
QwenImageControlNetInpaintPipeline,
854854
QwenImageControlNetPipeline,

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
2424
else:
2525
_import_structure["pipeline_prx"] = ["PRXPipeline"]
26+
_import_structure["pipeline_prx_pixel"] = ["PRXPixelPipeline"]
2627

2728
# Wrap T5GemmaEncoder to pass config.encoder (T5GemmaModuleConfig) instead of the
2829
# composite T5GemmaConfig, which lacks flat attributes expected by T5GemmaEncoder.__init__.
@@ -57,6 +58,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
5758
else:
5859
from .pipeline_output import PRXPipelineOutput
5960
from .pipeline_prx import PRXPipeline
61+
from .pipeline_prx_pixel import PRXPixelPipeline
6062

6163
else:
6264
import sys

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

Lines changed: 47 additions & 13 deletions

0 commit comments

Comments
 (0)