github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add support for PRX Pixel T2I

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-13 09:47:53 -06:00
parent a01c83073a
commit 6c3b82696e
9 changed files with 1464 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
extensions_built_in/diffusion_models/prx_pixel_t2i/src/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
# Everything diffusers does NOT ship (until PR #13928 lands) lives in src/:
# the vendored PRX transformer architecture and a minimal pixel-space sampler.
from .transformer_prx import PRXTransformer2DModel
from .pipeline import PRXPixelPipeline

154
extensions_built_in/diffusion_models/prx_pixel_t2i/src/pipeline.py Normal file
View File

@@ -0,0 +1,154 @@
"""Minimal preview sampler for the pixel-space PRX (PRXPixel) model.
ai-toolkit only uses this pipeline to render preview/sample images during
training (BaseModel.generate_images -> PRXPixelT2IModel.generate_single_image).
It does not need to be a diffusers DiffusionPipeline, and ai-toolkit always
encodes the prompts itself, so the pipeline only ever receives already-encoded
``PromptEmbeds`` (text features + attention mask), never raw text.
PRXPixel specifics this sampler bakes in (see ../prx_pixel_t2i.py for the why):
- **pixel space**: there is no VAE. The model's "latents" are the RGB image
itself in [-1, 1]; the final "decode" is just a clamp + uint8 cast.
- **x-prediction**: the transformer predicts the clean image x0, not the
flow-matching velocity. Each step converts it to velocity
``v = (x_t - x0) / t`` (t clamped for stability) before the scheduler step,
exactly like the diffusers PRXPixelPipeline.
- **noise_scale**: PRXPixel trains with a non-unit initial-noise std, so the
starting noise is ``randn * noise_scale`` (2.0 for the released model).
- **CFG** is applied on the x0 prediction (before the velocity conversion).
"""
from typing import List, Optional
import torch
from PIL import Image
from diffusers.utils.torch_utils import randn_tensor
# Minimum normalized timestep used when converting an x0 prediction to a flow
# velocity ``v = (x_t - x0) / t``. Mirrors the 0.05 clamp in the diffusers
# PRXPixelPipeline; without it the division blows up as t -> 0.
X_PRED_T_MIN = 0.05
class PRXPixelPipeline:
"""Lightweight pixel-space flow-matching sampler used for training previews."""
def __init__(self, model):
# ``model`` is the PRXPixelT2IModel (a BaseModel subclass), giving us
# access to model.transformer, model.decode_latents, device/dtype, the
# scheduler factory and the noise scale.
self.model = model
@property
def device(self):
return self.model.device_torch
def to(self, *args, **kwargs):
# BaseModel.generate_images may call pipeline.to(device); we manage
# devices through the model itself, so this is a no-op.
return self
def set_progress_bar_config(self, **kwargs):
# called by the sampler harness (inside a try/except, so optional)
pass
def _embeds_and_mask(self, embeds, device, dtype):
"""Pull (features, attention_mask) out of a PromptEmbeds onto device/dtype.
``text_embeds`` is (B, L, D). The mask is (B, L) bool (1 = real token);
it is kept as long for the transformer's boolean masking and never cast
to the model dtype.
"""
feats = embeds.text_embeds.to(device, dtype=dtype)
mask = getattr(embeds, "attention_mask", None)
if mask is not None:
mask = mask.to(device)
return feats, mask
@torch.no_grad()
def __call__(
self,
conditional_embeds, # PromptEmbeds: .text_embeds (B,L,D) + .attention_mask
unconditional_embeds, # PromptEmbeds or None (negative prompt)
height: int = 1024,
width: int = 1024,
num_inference_steps: int = 28,
guidance_scale: float = 5.0,
latents: Optional[torch.Tensor] = None, # pre-made noise, usually None
generator: Optional[
torch.Generator
] = None, # seeded RNG for reproducible samples
**kwargs,
) -> List[Image.Image]:
model = self.model
device = model.device_torch
dtype = model.torch_dtype
transformer = model.transformer
# Always sample with a FRESH scheduler -- the training scheduler is
# stateful and mutating it mid-training would corrupt the train step.
scheduler = model.get_train_scheduler()
scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = scheduler.timesteps # 1000 -> 0 scale
do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
# 1. starting noise -- pixel space, so channels = transformer.in_channels (3),
# spatial size = the requested pixels (no VAE downsample). Scaled by the
# model's noise_scale to match the learned flow-matching trajectory.
if latents is None:
shape = (1, transformer.in_channels, height, width)
latents = randn_tensor(
shape, generator=generator, device=device, dtype=torch.float32
)
latents = latents * model.noise_scale
latents = latents.to(device, dtype=torch.float32)
# 2. text features + masks
cond_feats, cond_mask = self._embeds_and_mask(conditional_embeds, device, dtype)
if do_cfg:
uncond_feats, uncond_mask = self._embeds_and_mask(
unconditional_embeds, device, dtype
)
# 3. denoising loop
for t in timesteps:
# scheduler timesteps are 0-1000; the transformer wants [0, 1]
t01 = (t / 1000.0).to(device).float().expand(latents.shape[0])
x0_cond = transformer(
hidden_states=latents.to(dtype),
timestep=t01,
encoder_hidden_states=cond_feats,
attention_mask=cond_mask,
return_dict=False,
)[0]
if do_cfg:
x0_uncond = transformer(
hidden_states=latents.to(dtype),
timestep=t01,
encoder_hidden_states=uncond_feats,
attention_mask=uncond_mask,
return_dict=False,
)[0]
# classifier-free guidance is applied on the x0 prediction
x0 = x0_uncond + guidance_scale * (x0_cond - x0_uncond)
else:
x0 = x0_cond
# convert the x0 (clean-image) prediction to the flow velocity the
# scheduler consumes: v = (x_t - x0) / t, t clamped for stability.
t_x = torch.clamp(t01.to(torch.float32), min=X_PRED_T_MIN).view(-1, 1, 1, 1)
v = (latents - x0.to(torch.float32)) / t_x
latents = scheduler.step(v, t, latents, return_dict=False)[0]
# 4. pixel space: the denoised latents ARE the image in [-1, 1].
# decode_latents is an identity (FakeVAE) but we route through it so
# any future latent normalization stays in one place.
images = model.decode_latents(latents, device=device, dtype=torch.float32)
images = images.float().clamp(-1.0, 1.0)
images = ((images + 1.0) * 127.5).round().to(torch.uint8)
images = images.permute(0, 2, 3, 1).cpu().numpy()
return [Image.fromarray(arr) for arr in images]

881
extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py Normal file
View File

@@ -0,0 +1,881 @@
# PRXTransformer2DModel -- vendored, self-contained copy of the Photoroom PRX
# transformer so ai-toolkit does not depend on an unmerged diffusers PR
# (https://github.com/huggingface/diffusers/pull/13928).
#
# This is the architecture as it appears on the ``prx-pixel`` branch, with two
# small edits so it runs against a stock installed diffusers:
# - the relative ``from ...`` / ``from ..`` imports are made absolute
# (``from diffusers...``) -- every building block used here ships in stock
# diffusers already, only the assembled PRX model is missing.
# - ``maybe_adjust_dtype_for_device`` (added in a newer diffusers) is inlined
# below so we don't depend on its presence.
#
# The pixel-space PRX-7B variant (Photoroom/prxpixel-t2i) sets
# ``in_channels=3`` (raw RGB, no VAE), ``bottleneck_size`` (two-layer patch
# projection) and ``resolution_embeds=True`` (resolution-aware modulation).
#
# Copyright 2025 The Photoroom and The HuggingFace Teams. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Optional
import torch
from torch import nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils import logging
from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.embeddings import get_timestep_embedding
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.normalization import RMSNorm
logger = logging.get_logger(__name__)
def maybe_adjust_dtype_for_device(
dtype: torch.dtype, device: torch.device
) -> torch.dtype:
r"""
MPS does not implement float64; fall back to float32 there. On every other device the requested dtype is returned
unchanged. Inlined from newer diffusers so this file works on older installs.
"""
if (
device is not None
and getattr(device, "type", None) == "mps"
and dtype == torch.float64
):
return torch.float32
return dtype
def get_image_ids(
batch_size: int, height: int, width: int, patch_size: int, device: torch.device
) -> torch.Tensor:
r"""
Generates 2D patch coordinate indices for a batch of images.
Args:
batch_size (`int`):
Number of images in the batch.
height (`int`):
Height of the input images (in pixels).
width (`int`):
Width of the input images (in pixels).
patch_size (`int`):
Size of the square patches that the image is divided into.
device (`torch.device`):
The device on which to create the tensor.
Returns:
`torch.Tensor`:
Tensor of shape `(batch_size, num_patches, 2)` containing the (row, col) coordinates of each patch in the
image grid.
"""
img_ids = torch.zeros(height // patch_size, width // patch_size, 2, device=device)
img_ids[..., 0] = torch.arange(height // patch_size, device=device)[:, None]
img_ids[..., 1] = torch.arange(width // patch_size, device=device)[None, :]
return (
img_ids.reshape((height // patch_size) * (width // patch_size), 2)
.unsqueeze(0)
.repeat(batch_size, 1, 1)
)
def apply_rope(xq: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
r"""
Applies rotary positional embeddings (RoPE) to a query tensor.
Args:
xq (`torch.Tensor`):
Input tensor of shape `(..., dim)` representing the queries.
freqs_cis (`torch.Tensor`):
Precomputed rotary frequency components of shape `(..., dim/2, 2)` containing cosine and sine pairs.
Returns:
`torch.Tensor`:
Tensor of the same shape as `xq` with rotary embeddings applied.
"""
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
# Ensure freqs_cis is on the same device as queries to avoid device mismatches with offloading
freqs_cis = freqs_cis.to(device=xq.device, dtype=xq_.dtype)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
return xq_out.reshape(*xq.shape).type_as(xq)
class PRXAttnProcessor2_0:
r"""
Processor for implementing PRX-style attention with multi-source tokens and RoPE. Supports multiple attention
backends (Flash Attention, Sage Attention, etc.) via dispatch_attention_fn.
"""
_attention_backend = None
_parallel_config = None
def __init__(self):
if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
raise ImportError(
"PRXAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0."
)
def __call__(
self,
attn: "PRXAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
"""
Apply PRX attention using PRXAttention module.
Args:
attn: PRXAttention module containing projection layers
hidden_states: Image tokens [B, L_img, D]
encoder_hidden_states: Text tokens [B, L_txt, D]
attention_mask: Boolean mask for text tokens [B, L_txt]
image_rotary_emb: Rotary positional embeddings [B, 1, L_img, head_dim//2, 2, 2]
"""
if encoder_hidden_states is None:
raise ValueError(
"PRXAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens."
)
# Project image tokens to Q, K, V
img_qkv = attn.img_qkv_proj(hidden_states)
B, L_img, _ = img_qkv.shape
img_qkv = img_qkv.reshape(B, L_img, 3, attn.heads, attn.head_dim)
img_qkv = img_qkv.permute(2, 0, 3, 1, 4) # [3, B, H, L_img, D]
img_q, img_k, img_v = img_qkv[0], img_qkv[1], img_qkv[2]
# Apply QK normalization to image tokens
img_q = attn.norm_q(img_q)
img_k = attn.norm_k(img_k)
# Project text tokens to K, V
txt_kv = attn.txt_kv_proj(encoder_hidden_states)
B, L_txt, _ = txt_kv.shape
txt_kv = txt_kv.reshape(B, L_txt, 2, attn.heads, attn.head_dim)
txt_kv = txt_kv.permute(2, 0, 3, 1, 4) # [2, B, H, L_txt, D]
txt_k, txt_v = txt_kv[0], txt_kv[1]
# Apply K normalization to text tokens
txt_k = attn.norm_added_k(txt_k)
# Apply RoPE to image queries and keys
if image_rotary_emb is not None:
img_q = apply_rope(img_q, image_rotary_emb)
img_k = apply_rope(img_k, image_rotary_emb)
# Concatenate text and image keys/values
k = torch.cat((txt_k, img_k), dim=2) # [B, H, L_txt + L_img, D]
v = torch.cat((txt_v, img_v), dim=2) # [B, H, L_txt + L_img, D]
# Build attention mask if provided
attn_mask_tensor = None
if attention_mask is not None:
bs, _, l_img, _ = img_q.shape
l_txt = txt_k.shape[2]
if attention_mask.dim() != 2:
raise ValueError(
f"Unsupported attention_mask shape: {attention_mask.shape}"
)
if attention_mask.shape[-1] != l_txt:
raise ValueError(
f"attention_mask last dim {attention_mask.shape[-1]} must equal text length {l_txt}"
)
device = img_q.device
ones_img = torch.ones((bs, l_img), dtype=torch.bool, device=device)
attention_mask = attention_mask.to(device=device, dtype=torch.bool)
joint_mask = torch.cat([attention_mask, ones_img], dim=-1)
attn_mask_tensor = joint_mask[:, None, None, :].expand(
-1, attn.heads, l_img, -1
)
# Apply attention using dispatch_attention_fn for backend support
# Reshape to match dispatch_attention_fn expectations: [B, L, H, D]
query = img_q.transpose(1, 2) # [B, L_img, H, D]
key = k.transpose(1, 2) # [B, L_txt + L_img, H, D]
value = v.transpose(1, 2) # [B, L_txt + L_img, H, D]
attn_output = dispatch_attention_fn(
query,
key,
value,
attn_mask=attn_mask_tensor,
backend=self._attention_backend,
parallel_config=self._parallel_config,
)
# Reshape from [B, L_img, H, D] to [B, L_img, H*D]
batch_size, seq_len, num_heads, head_dim = attn_output.shape
attn_output = attn_output.reshape(batch_size, seq_len, num_heads * head_dim)
# Apply output projection
attn_output = attn.to_out[0](attn_output)
if len(attn.to_out) > 1:
attn_output = attn.to_out[1](attn_output) # dropout if present
return attn_output
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]
def __init__(
self,
query_dim: int,
heads: int = 8,
dim_head: int = 64,
bias: bool = False,
out_bias: bool = False,
eps: float = 1e-6,
processor=None,
):
super().__init__()
self.heads = heads
self.head_dim = dim_head
self.inner_dim = dim_head * heads
self.query_dim = query_dim
self.img_qkv_proj = nn.Linear(query_dim, query_dim * 3, bias=bias)
self.norm_q = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True)
self.norm_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True)
self.txt_kv_proj = nn.Linear(query_dim, query_dim * 2, bias=bias)
self.norm_added_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True)
self.to_out = nn.ModuleList([])
self.to_out.append(nn.Linear(self.inner_dim, query_dim, bias=out_bias))
self.to_out.append(nn.Dropout(0.0))
if processor is None:
processor = self._default_processor_cls()
self.set_processor(processor)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
image_rotary_emb: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
return self.processor(
self,
hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
image_rotary_emb=image_rotary_emb,
**kwargs,
)
# inspired from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
class PRXEmbedND(nn.Module):
r"""
N-dimensional rotary positional embedding.
This module creates rotary embeddings (RoPE) across multiple axes, where each axis can have its own embedding
dimension. The embeddings are combined and returned as a single tensor
Args:
dim (int):
Base embedding dimension (must be even).
theta (int):
Scaling factor that controls the frequency spectrum of the rotary embeddings.
axes_dim (list[int]):
list of embedding dimensions for each axis (each must be even).
"""
def __init__(self, dim: int, theta: int, axes_dim: list[int]):
super().__init__()
self.dim = dim
self.theta = theta
self.axes_dim = axes_dim
def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
assert dim % 2 == 0
dtype = maybe_adjust_dtype_for_device(torch.float64, pos.device)
scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim
omega = 1.0 / (theta**scale)
out = pos.unsqueeze(-1) * omega.unsqueeze(0)
out = torch.stack(
[torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1
)
# Native PyTorch equivalent of: Rearrange("b n d (i j) -> b n d i j", i=2, j=2)
# out shape: (b, n, d, 4) -> reshape to (b, n, d, 2, 2)
out = out.reshape(*out.shape[:-1], 2, 2)
return out.float()
def forward(self, ids: torch.Tensor) -> torch.Tensor:
n_axes = ids.shape[-1]
emb = torch.cat(
[
self.rope(ids[:, :, i], self.axes_dim[i], self.theta)
for i in range(n_axes)
],
dim=-3,
)
return emb.unsqueeze(1)
class MLPEmbedder(nn.Module):
r"""
A simple 2-layer MLP used for embedding inputs.
Args:
in_dim (`int`):
Dimensionality of the input features.
hidden_dim (`int`):
Dimensionality of the hidden and output embedding space.
Returns:
`torch.Tensor`:
Tensor of shape `(..., hidden_dim)` containing the embedded representations.
"""
def __init__(self, in_dim: int, hidden_dim: int):
super().__init__()
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
self.silu = nn.SiLU()
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.out_layer(self.silu(self.in_layer(x)))
class PRXResolutionEmbedder(nn.Module):
r"""
Embeds the spatial resolution `(height, width)` of the latent into a vector that is added to the timestep
embedding, so the model can condition its modulation on the generation resolution.
A sinusoidal embedding of dimension 128 is built for the height and the width separately and concatenated into a
256-dim vector, which is then projected to `hidden_size` by a 2-layer MLP. This matches the `"vec"` mode of the
resolution-aware conditioning used during PRX-7B training.
Args:
hidden_size (`int`):
Dimension of the output embedding (must match the timestep embedding dimension).
max_period (`int`, *optional*, defaults to 10000):
Maximum frequency period for the sinusoidal resolution embedding.
"""
def __init__(self, hidden_size: int, max_period: int = 10000):
super().__init__()
self.max_period = max_period
self.mlp = MLPEmbedder(in_dim=256, hidden_dim=hidden_size)
def forward(
self, height: torch.Tensor, width: torch.Tensor, dtype: torch.dtype
) -> torch.Tensor:
h_emb = get_timestep_embedding(
timesteps=height,
embedding_dim=128,
max_period=self.max_period,
scale=1.0,
flip_sin_to_cos=True,
downscale_freq_shift=0.0,
)
w_emb = get_timestep_embedding(
timesteps=width,
embedding_dim=128,
max_period=self.max_period,
scale=1.0,
flip_sin_to_cos=True,
downscale_freq_shift=0.0,
)
hw_emb = torch.cat([h_emb, w_emb], dim=-1).to(dtype)
return self.mlp(hw_emb)
class Modulation(nn.Module):
r"""
Modulation network that generates scale, shift, and gating parameters.
Given an input vector, the module projects it through a linear layer to produce six chunks, which are grouped into
two tuples `(shift, scale, gate)`.
Args:
dim (`int`):
Dimensionality of the input vector. The output will have `6 * dim` features internally.
Returns:
((`torch.Tensor`, `torch.Tensor`, `torch.Tensor`), (`torch.Tensor`, `torch.Tensor`, `torch.Tensor`)):
Two tuples `(shift, scale, gate)`.
"""
def __init__(self, dim: int):
super().__init__()
self.lin = nn.Linear(dim, 6 * dim, bias=True)
nn.init.constant_(self.lin.weight, 0)
nn.init.constant_(self.lin.bias, 0)
def forward(
self, vec: torch.Tensor
) -> tuple[
tuple[torch.Tensor, torch.Tensor, torch.Tensor],
tuple[torch.Tensor, torch.Tensor, torch.Tensor],
]:
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(6, dim=-1)
return tuple(out[:3]), tuple(out[3:])
class PRXBlock(nn.Module):
r"""
Multimodal transformer block with text-image cross-attention, modulation, and MLP.
Args:
hidden_size (`int`):
Dimension of the hidden representations.
num_heads (`int`):
Number of attention heads.
mlp_ratio (`float`, *optional*, defaults to 4.0):
Expansion ratio for the hidden dimension inside the MLP.
qk_scale (`float`, *optional*):
Scale factor for queries and keys. If not provided, defaults to ``head_dim**-0.5``.
"""
def __init__(
self,
hidden_size: int,
num_heads: int,
mlp_ratio: float = 4.0,
qk_scale: float | None = None,
):
super().__init__()
self.hidden_dim = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.scale = qk_scale or self.head_dim**-0.5
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
self.hidden_size = hidden_size
# Pre-attention normalization for image tokens
self.img_pre_norm = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6
)
# PRXAttention module with built-in projections and norms
self.attention = PRXAttention(
query_dim=hidden_size,
heads=num_heads,
dim_head=self.head_dim,
bias=False,
out_bias=False,
eps=1e-6,
processor=PRXAttnProcessor2_0(),
)
# mlp
self.post_attention_layernorm = nn.LayerNorm(
hidden_size, elementwise_affine=False, eps=1e-6
)
self.gate_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False)
self.up_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False)
self.down_proj = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False)
self.mlp_act = nn.GELU(approximate="tanh")
self.modulation = Modulation(hidden_size)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
image_rotary_emb: torch.Tensor,
attention_mask: torch.Tensor | None = None,
**kwargs: dict[str, Any],
) -> torch.Tensor:
r"""
Runs modulation-gated cross-attention and MLP, with residual connections.
Args:
hidden_states (`torch.Tensor`):
Image tokens of shape `(B, L_img, hidden_size)`.
encoder_hidden_states (`torch.Tensor`):
Text tokens of shape `(B, L_txt, hidden_size)`.
temb (`torch.Tensor`):
Conditioning vector used by `Modulation` to produce scale/shift/gates, shape `(B, hidden_size)` (or
broadcastable).
image_rotary_emb (`torch.Tensor`):
Rotary positional embeddings applied inside attention.
attention_mask (`torch.Tensor`, *optional*):
Boolean mask for text tokens of shape `(B, L_txt)`, where `0` marks padding.
Returns:
`torch.Tensor`:
Updated image tokens of shape `(B, L_img, hidden_size)`.
"""
mod_attn, mod_mlp = self.modulation(temb)
attn_shift, attn_scale, attn_gate = mod_attn
mlp_shift, mlp_scale, mlp_gate = mod_mlp
hidden_states_mod = (1 + attn_scale) * self.img_pre_norm(
hidden_states
) + attn_shift
attn_out = self.attention(
hidden_states=hidden_states_mod,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
image_rotary_emb=image_rotary_emb,
)
hidden_states = hidden_states + attn_gate * attn_out
x = (1 + mlp_scale) * self.post_attention_layernorm(hidden_states) + mlp_shift
hidden_states = hidden_states + mlp_gate * (
self.down_proj(self.mlp_act(self.gate_proj(x)) * self.up_proj(x))
)
return hidden_states
class FinalLayer(nn.Module):
r"""
Final projection layer with adaptive LayerNorm modulation.
This layer applies a normalized and modulated transformation to input tokens and projects them into patch-level
outputs.
Args:
hidden_size (`int`):
Dimensionality of the input tokens.
patch_size (`int`):
Size of the square image patches.
out_channels (`int`):
Number of output channels per pixel (e.g. RGB = 3).
"""
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(
hidden_size, patch_size * patch_size * out_channels, bias=True
)
self.adaLN_modulation = nn.Sequential(
nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
)
def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
x = self.linear(x)
return x
def img2seq(img: torch.Tensor, patch_size: int) -> torch.Tensor:
r"""
Flattens an image tensor into a sequence of non-overlapping patches.
Args:
img (`torch.Tensor`):
Input image tensor of shape `(B, C, H, W)`.
patch_size (`int`):
Size of each square patch. Must evenly divide both `H` and `W`.
Returns:
`torch.Tensor`:
Flattened patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W
// patch_size)` is the number of patches.
"""
b, c, h, w = img.shape
p = patch_size
# Reshape to (B, C, H//p, p, W//p, p) separating grid and patch dimensions
img = img.reshape(b, c, h // p, p, w // p, p)
# Permute to (B, H//p, W//p, C, p, p) using einsum
# n=batch, c=channels, h=grid_height, p=patch_height, w=grid_width, q=patch_width
img = torch.einsum("nchpwq->nhwcpq", img)
# Flatten to (B, L, C * p * p)
img = img.reshape(b, -1, c * p * p)
return img
def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Tensor:
r"""
Reconstructs an image tensor from a sequence of patches (inverse of `img2seq`).
Args:
seq (`torch.Tensor`):
Patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W //
patch_size)`.
patch_size (`int`):
Size of each square patch.
shape (`tuple` or `torch.Tensor`):
The original image spatial shape `(H, W)`. If a tensor is provided, the first two values are interpreted as
height and width.
Returns:
`torch.Tensor`:
Reconstructed image tensor of shape `(B, C, H, W)`.
"""
if isinstance(shape, tuple):
h, w = shape[-2:]
elif isinstance(shape, torch.Tensor):
h, w = (int(shape[0]), int(shape[1]))
else:
raise NotImplementedError(f"shape type {type(shape)} not supported")
b, l, d = seq.shape
p = patch_size
c = d // (p * p)
# Reshape back to grid structure: (B, H//p, W//p, C, p, p)
seq = seq.reshape(b, h // p, w // p, c, p, p)
# Permute back to image layout: (B, C, H//p, p, W//p, p)
# n=batch, h=grid_height, w=grid_width, c=channels, p=patch_height, q=patch_width
seq = torch.einsum("nhwcpq->nchpwq", seq)
# Final reshape to (B, C, H, W)
seq = seq.reshape(b, c, h, w)
return seq
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. The pixel-space PRX uses `3` (raw RGB, no VAE).
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.
bottleneck_size (`int`, *optional*):
If set, the image patch projection (`img_in`) uses a two-layer bottleneck (`patch_dim -> bottleneck_size ->
hidden_size`) instead of a single linear layer. Used by the pixel-space PRX-7B variant where the patch
dimension is large.
resolution_embeds (`bool`, *optional*, defaults to `False`):
Whether to condition the timestep modulation on the latent resolution `(H, W)` via a
`PRXResolutionEmbedder`. Used by the PRX-7B variant.
"""
config_name = "config.json"
_supports_gradient_checkpointing = True
@register_to_config
def __init__(
self,
in_channels: int = 16,
patch_size: int = 2,
context_in_dim: int = 2304,
hidden_size: int = 1792,
mlp_ratio: float = 3.5,
num_heads: int = 28,
depth: int = 16,
axes_dim: list = None,
theta: int = 10000,
time_factor: float = 1000.0,
time_max_period: int = 10000,
bottleneck_size: Optional[int] = None,
resolution_embeds: bool = False,
):
super().__init__()
if axes_dim is None:
axes_dim = [32, 32]
# Store parameters directly
self.in_channels = in_channels
self.patch_size = patch_size
self.out_channels = self.in_channels * self.patch_size**2
self.time_factor = time_factor
self.time_max_period = time_max_period
if hidden_size % num_heads != 0:
raise ValueError(
f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}"
)
pe_dim = hidden_size // num_heads
if sum(axes_dim) != pe_dim:
raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
self.hidden_size = hidden_size
self.num_heads = num_heads
self.pe_embedder = PRXEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
patch_dim = self.in_channels * self.patch_size**2
if bottleneck_size is not None:
# Two-layer bottleneck projection (used by pixel-space PRX where the patch dimension is large).
self.img_in = nn.Sequential(
nn.Linear(patch_dim, bottleneck_size, bias=True),
nn.Linear(bottleneck_size, self.hidden_size, bias=True),
)
else:
self.img_in = nn.Linear(patch_dim, self.hidden_size, bias=True)
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
self.resolution_embedder = (
PRXResolutionEmbedder(self.hidden_size, max_period=time_max_period)
if resolution_embeds
else None
)
self.blocks = nn.ModuleList(
[
PRXBlock(
self.hidden_size,
self.num_heads,
mlp_ratio=mlp_ratio,
)
for i in range(depth)
]
)
self.final_layer = FinalLayer(self.hidden_size, 1, self.out_channels)
self.gradient_checkpointing = False
def _compute_timestep_embedding(
self, timestep: torch.Tensor, dtype: torch.dtype
) -> torch.Tensor:
return self.time_in(
get_timestep_embedding(
timesteps=timestep,
embedding_dim=256,
max_period=self.time_max_period,
scale=self.time_factor,
flip_sin_to_cos=True, # Match original cos, sin order
downscale_freq_shift=0.0,
).to(dtype)
)
def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
attention_kwargs: dict[str, Any] | None = None,
return_dict: bool = True,
) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput:
r"""
Forward pass of the PRXTransformer2DModel.
The latent image is split into patch tokens, combined with text conditioning, and processed through a stack of
transformer blocks modulated by the timestep. The output is reconstructed into the latent image space.
Args:
hidden_states (`torch.Tensor`):
Input latent image tensor of shape `(B, C, H, W)`.
timestep (`torch.Tensor`):
Timestep tensor of shape `(B,)` or `(1,)` on a `[0, 1]` scale, used for temporal conditioning.
encoder_hidden_states (`torch.Tensor`):
Text conditioning tensor of shape `(B, L_txt, context_in_dim)`.
attention_mask (`torch.Tensor`, *optional*):
Boolean mask of shape `(B, L_txt)`, where `0` marks padding in the text sequence.
Returns:
`Transformer2DModelOutput` if `return_dict=True`, otherwise a tuple `(sample,)` with the predicted clean
image of shape `(B, C, H, W)`.
"""
# Process text conditioning
txt = self.txt_in(encoder_hidden_states)
# Convert image to sequence and embed
img = img2seq(hidden_states, self.patch_size)
img = self.img_in(img)
# Generate positional embeddings
bs, _, h, w = hidden_states.shape
img_ids = get_image_ids(
bs, h, w, patch_size=self.patch_size, device=hidden_states.device
)
pe = self.pe_embedder(img_ids)
# Compute time embedding
vec = self._compute_timestep_embedding(timestep, dtype=img.dtype)
# Add resolution conditioning (PRX-7B "vec" mode): embed the latent (H, W) and add it to the timestep vector
# so every block's modulation is resolution-aware.
if self.resolution_embedder is not None:
height = torch.full(
(bs,), h, device=hidden_states.device, dtype=torch.float32
)
width = torch.full(
(bs,), w, device=hidden_states.device, dtype=torch.float32
)
vec = vec + self.resolution_embedder(height, width, dtype=vec.dtype)
# Apply transformer blocks
for block in self.blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
img = self._gradient_checkpointing_func(
block.__call__,
img,
txt,
vec,
pe,
attention_mask,
)
else:
img = block(
hidden_states=img,
encoder_hidden_states=txt,
temb=vec,
image_rotary_emb=pe,
attention_mask=attention_mask,
)
# Final layer and convert back to image
img = self.final_layer(img, vec)
output = seq2img(img, self.patch_size, hidden_states.shape)
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)