Add support for PRX Pixel T2I

README.md

@@ -27,6 +27,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [baidu/ERNIE-Image](https://huggingface.co/baidu/ERNIE-Image) (ERNIE-Image)
 - [NucleusAI/Nucleus-Image](https://huggingface.co/NucleusAI/Nucleus-Image) (Nucleus-Image)
 - [HiDream-ai/HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) (HiDream O1)
+- [Photoroom/prxpixel-t2i](https://huggingface.co/Photoroom/prxpixel-t2i) (PRXPixel)
 
 ### Instruction / Edit
 - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev)

extensions_built_in/diffusion_models/__init__.py

@@ -14,6 +14,7 @@ from .nucleus_image import NucleusImageModel
 from .hidream.hidream_o1_model import HidreamO1Model
 from .z_image.z_image_l2p_model import ZImageL2PModel
 from .ideogram4 import Ideogram4Model
+from .prx_pixel_t2i import PRXPixelT2IModel
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
@@ -42,4 +43,5 @@ AI_TOOLKIT_MODELS = [
     HidreamO1Model,
     ZImageL2PModel,
     Ideogram4Model,
+    PRXPixelT2IModel,
 ]

extensions_built_in/diffusion_models/prx_pixel_t2i/README.md

@@ -0,0 +1,53 @@
+# PRXPixel (Photoroom PRX-7B, pixel-space text-to-image)
+
+ai-toolkit integration for [`Photoroom/prxpixel-t2i`](https://huggingface.co/Photoroom/prxpixel-t2i),
+a ~7B pixel-space diffusion transformer.
+
+It is implemented **from scratch** so ai-toolkit does not depend on the (still
+unmerged) diffusers PR [huggingface/diffusers#13928](https://github.com/huggingface/diffusers/pull/13928):
+the transformer is vendored in [src/transformer_prx.py](src/transformer_prx.py)
+and a minimal preview sampler lives in [src/pipeline.py](src/pipeline.py).
+
+## What makes this model unusual
+
+PRXPixel differs from a typical latent flow-matching model in three ways, each
+handled in [prx_pixel_t2i.py](prx_pixel_t2i.py):
+
+| Property | What it means | How it's handled |
+|---|---|---|
+| **Pixel space** | No VAE — the transformer denoises raw RGB (`in_channels=3`, `patch_size=16`) | A `FakeVAE` (identity, scaling 1) so encode/decode are no-ops; "latents" are the image in `[-1, 1]` |
+| **x-prediction** | The model predicts the clean image `x0`, not the flow velocity | `get_noise_prediction` returns `x0`; `get_loss_target` is the clean latents. The `x0 → velocity` conversion only happens at sampling time |
+| **noise_scale = 2.0** | Trains/samples from `randn * 2.0`, not unit noise | `get_latent_noise_from_latents` scales the training noise; the pipeline scales the starting noise |
+
+Text is encoded by the Qwen3-VL text tower (`Qwen3VLTextModel`, hidden size
+2048 → the transformer's `context_in_dim`), padded to 256 tokens.
+
+The x-prediction objective follows *"Back to Basics: Let Denoising Generative
+Models Denoise"* (https://arxiv.org/abs/2511.13720).
+
+## Architecture (released checkpoint)
+
+`depth=24`, `hidden_size=3584`, `num_heads=28`, `mlp_ratio=3.5`,
+`in_channels=3`, `patch_size=16`, `context_in_dim=2048`, `bottleneck_size=768`,
+`axes_dim=[64, 64]`, `resolution_embeds=True`, flow-matching scheduler with
+`shift=3.0`.
+
+## Train it
+
+```yaml
+model:
+  arch: "prx_pixel"
+  name_or_path: "/path/to/prxpixel-t2i"   # diffusers folder: transformer/,
+                                          # text_encoder/, tokenizer/, scheduler/
+  quantize: true        # optional: qfloat8 the transformer
+  quantize_te: true     # optional: qfloat8 the Qwen3-VL text encoder
+train:
+  gradient_checkpointing: true
+sample:
+  guidance_scale: 5.0
+  sample_steps: 28
+```
+
+Datasets bucket to multiples of 16px (`vae_scale_factor * patch_size`).
+See [../example_model/README.md](../example_model/README.md) for the generic
+lifecycle, registration and LoRA conventions.

extensions_built_in/diffusion_models/prx_pixel_t2i/__init__.py

@@ -0,0 +1,3 @@
+from .prx_pixel_t2i import PRXPixelT2IModel
+
+__all__ = ["PRXPixelT2IModel"]

extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py

@@ -0,0 +1,345 @@
+"""PRXPixelT2IModel -- the Photoroom pixel-space PRX-7B text-to-image model
+(https://huggingface.co/Photoroom/prxpixel-t2i) wired into ai-toolkit.
+
+This is implemented from scratch (the diffusers support is an unmerged PR,
+https://github.com/huggingface/diffusers/pull/13928) so ai-toolkit does not
+depend on it: the transformer architecture is vendored in ``src/transformer_prx.py``
+and a minimal sampler lives in ``src/pipeline.py``.
+
+What makes PRXPixel unusual (and what each override below is doing about it):
+
+  - **Pixel space, no VAE.** The transformer denoises raw RGB directly
+    (``in_channels=3``, ``patch_size=16``). We use a ``FakeVAE`` (identity,
+    scaling_factor=1) so BaseModel's encode_images/decode_latents become no-ops
+    and the "latents" everywhere in the toolkit are just the image in [-1, 1].
+    Same trick as ``../chroma/chroma_radiance_model.py`` and
+    ``extensions/z_image_pixel``.
+
+  - **x-prediction.** The model predicts the CLEAN image x0, not the
+    flow-matching velocity. ai-toolkit's MSE compares ``get_noise_prediction``
+    against ``get_loss_target``; we set BOTH to the x0 space (prediction = the
+    model's x0 output, target = the clean latents), which is PRXPixel's native
+    training objective ("Back to Basics: Let Denoising Generative Models
+    Denoise", https://arxiv.org/abs/2511.13720). The x0->velocity conversion
+    only happens at sampling time, inside the pipeline.
+
+  - **noise_scale = 2.0.** PRXPixel trains with a non-unit initial-noise std,
+    so the noise mixed into the latents (training) and the starting noise
+    (sampling) are ``randn * noise_scale``. We override
+    ``get_latent_noise_from_latents`` for the training side; the pipeline
+    handles the sampling side.
+
+  - **Qwen3-VL text tower.** Prompts are encoded by ``Qwen3VLTextModel``
+    (hidden size 2048 -> the transformer's ``context_in_dim``). We keep the
+    per-token ``last_hidden_state`` plus an attention mask.
+
+See ../example_model/README.md for the generic lifecycle/registration guide.
+"""
+
+import os
+from typing import List, Optional
+
+import torch
+import yaml
+
+from transformers import AutoTokenizer, Qwen3VLTextModel
+from optimum.quanto import freeze
+
+from toolkit.accelerator import unwrap_model
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.models.base_model import BaseModel
+from toolkit.models.FakeVAE import FakeVAE
+from toolkit.prompt_utils import PromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+from toolkit.train_tools import apply_noise_offset
+from toolkit.util.quantize import quantize, get_qtype, quantize_model
+
+from .src.transformer_prx import PRXTransformer2DModel
+from .src.pipeline import PRXPixelPipeline
+
+
+# Flow-matching scheduler config, matching the released model's
+# scheduler/scheduler_config.json (shift 3.0, 1000 train timesteps).
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "use_dynamic_shifting": False,
+    "shift": 3.0,
+}
+
+# Number of text tokens PRXPixel was trained with (the Qwen tokenizer's own
+# model_max_length is far larger). Matches PRX_PIXEL_DEFAULT_MAX_TOKENS.
+PROMPT_MAX_TOKENS = 256
+# Initial-noise std PRXPixel trains/samples with.
+NOISE_SCALE = 2.0
+
+
+class PRXPixelT2IModel(BaseModel):
+    # ``model.arch: "prx_pixel"`` in the training config YAML selects this class.
+    arch = "prx_pixel"
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        # Matched against type(module).__name__ to place LoRA layers.
+        self.target_lora_modules = ["PRXTransformer2DModel"]
+
+        # used by our overrides below
+        self.patch_size = 16
+        self.vae_scale_factor = 1  # pixel space: no VAE downsampling
+        self.max_text_length = PROMPT_MAX_TOKENS
+        self.noise_scale = NOISE_SCALE
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    def get_bucket_divisibility(self):
+        # pixels must be divisible by patch_size (no VAE downsample): 1 * 16 = 16
+        return self.vae_scale_factor * self.patch_size
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def load_model(self):
+        dtype = self.torch_dtype
+        self.print_and_status_update("Loading PRXPixel model")
+        # Expected diffusers-style layout under name_or_path:
+        #   transformer/ (PRXTransformer2DModel)  text_encoder/ (Qwen3VLTextModel)
+        #   tokenizer/                            scheduler/
+        model_path = self.model_config.name_or_path
+
+        # --- transformer (vendored PRXTransformer2DModel) ---
+        self.print_and_status_update("Loading transformer")
+        # from_pretrained reads config.json (bottleneck_size, resolution_embeds,
+        # in_channels=3, ...) and the safetensors in one shot.
+        transformer = PRXTransformer2DModel.from_pretrained(
+            model_path, subfolder="transformer", torch_dtype=dtype
+        )
+        transformer.to(dtype=dtype)
+        flush()
+
+        if self.model_config.quantize:
+            self.print_and_status_update("Quantizing transformer")
+            quantize_model(self, transformer)
+            flush()
+
+        if self.model_config.low_vram:
+            transformer.to("cpu")
+        else:
+            transformer.to(self.device_torch, dtype=dtype)
+        flush()
+
+        # --- text encoder + tokenizer (Qwen3-VL text tower) ---
+        self.print_and_status_update("Loading text encoder")
+        tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+        text_encoder = Qwen3VLTextModel.from_pretrained(
+            model_path, subfolder="text_encoder", torch_dtype=dtype
+        )
+        text_encoder.to(self.te_device_torch)
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        flush()
+
+        if self.model_config.quantize_te:
+            self.print_and_status_update("Quantizing text encoder")
+            quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
+            freeze(text_encoder)
+            flush()
+
+        # --- "VAE": identity, since PRXPixel is pixel space ---
+        self.print_and_status_update("Preparing pixel-space VAE (identity)")
+        vae = FakeVAE(scaling_factor=1.0)
+        vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype)
+
+        # --- scheduler + store everything ---
+        self.noise_scheduler = PRXPixelT2IModel.get_train_scheduler()
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.model = transformer  # aliased as self.transformer / self.unet
+        self.pipeline = PRXPixelPipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Sampling (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return PRXPixelPipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: PRXPixelPipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: PromptEmbeds,
+        unconditional_embeds: PromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        # snap requested size to the model's divisibility (16px)
+        sc = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // sc * sc)
+        gen_config.height = int(gen_config.height // sc * sc)
+
+        img = pipeline(
+            conditional_embeds=conditional_embeds,
+            unconditional_embeds=unconditional_embeds,
+            height=gen_config.height,
+            width=gen_config.width,
+            num_inference_steps=gen_config.num_inference_steps,
+            guidance_scale=gen_config.guidance_scale,
+            latents=gen_config.latents,
+            generator=generator,
+        )[0]
+        return img
+
+    # ------------------------------------------------------------------
+    # Training hooks
+    # ------------------------------------------------------------------
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,
+        timestep: torch.Tensor,  # 0..1000 scale, 1000 = pure noise
+        text_embeddings: PromptEmbeds,
+        **kwargs,
+    ):
+        """Forward pass. Returns the model's predicted CLEAN image x0.
+
+        PRXPixel is an x-prediction model, so the raw transformer output is the
+        prediction we compare against ``get_loss_target`` (the clean latents).
+        No velocity conversion happens here -- that is a sampling-time concern.
+        """
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        # toolkit timestep (0..1000) -> PRX flow time in [0, 1].
+        t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0
+
+        feats = text_embeddings.text_embeds.to(self.device_torch, self.torch_dtype)
+        mask = getattr(text_embeddings, "attention_mask", None)
+        if mask is not None:
+            mask = mask.to(self.device_torch)
+
+        x0_pred = self.model(
+            hidden_states=latent_model_input.to(self.device_torch, self.torch_dtype),
+            timestep=t01,
+            encoder_hidden_states=feats,
+            attention_mask=mask,
+            return_dict=False,
+        )[0]
+        return x0_pred
+
+    def get_prompt_embeds(self, prompt) -> PromptEmbeds:
+        """Encode prompt text with the Qwen3-VL text tower.
+
+        Returns a PromptEmbeds whose ``text_embeds`` is (B, L, 2048) and whose
+        ``attention_mask`` is (B, L). Prompts are padded to a fixed length
+        (PROMPT_MAX_TOKENS) so cached embeds are concatenatable for CFG.
+        """
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.max_text_length,
+            truncation=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        input_ids = text_inputs.input_ids.to(self.text_encoder.device)
+        attention_mask = text_inputs.attention_mask.to(self.text_encoder.device)
+
+        with torch.no_grad():
+            output = self.text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        embeds = output["last_hidden_state"].to(self.torch_dtype)
+
+        pe = PromptEmbeds(embeds)
+        # keep the mask as bool/long; it must not be dtype-cast with the embeds
+        pe.attention_mask = attention_mask.bool()
+        return pe
+
+    def get_loss_target(self, *args, **kwargs):
+        """x-prediction target: the clean image (x0).
+
+        PRXPixel predicts x0 directly, so the MSE target is simply the clean
+        latents (the pixel image in [-1, 1]), not the flow velocity.
+        """
+        batch = kwargs.get("batch")
+        return batch.latents.detach()
+
+    def get_latent_noise_from_latents(self, latents: torch.Tensor, noise_offset=0.0):
+        """Noise for the forward flow, scaled by the model's noise_scale.
+
+        PRXPixel trains with a non-unit initial-noise std: the noise mixed into
+        the latents is ``randn * noise_scale``. The scheduler's add_noise then
+        forms ``x_t = (1 - t) * clean + t * noise``.
+        """
+        noise = torch.randn_like(latents) * self.noise_scale
+        if noise_offset is not None and noise_offset != 0.0:
+            noise = apply_noise_offset(noise, noise_offset)
+        return noise
+
+    def condition_noisy_latents(self, latents: torch.Tensor, batch):
+        # plain text-to-image: nothing to inject
+        return latents
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    def save_model(self, output_path, meta, save_dtype):
+        """Full fine-tune save: write the transformer back in diffusers layout."""
+        transformer: PRXTransformer2DModel = unwrap_model(self.model)
+        transformer.save_pretrained(
+            os.path.join(output_path, "transformer"),
+            safe_serialization=True,
+        )
+        with open(os.path.join(output_path, "aitk_meta.yaml"), "w") as f:
+            yaml.dump(meta, f)
+
+    def get_base_model_version(self):
+        return "prx_pixel"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return ["blocks"]
+
+    def convert_lora_weights_before_save(self, state_dict):
+        return {
+            k.replace("transformer.", "diffusion_model."): v
+            for k, v in state_dict.items()
+        }
+
+    def convert_lora_weights_before_load(self, state_dict):
+        return {
+            k.replace("diffusion_model.", "transformer."): v
+            for k, v in state_dict.items()
+        }

extensions_built_in/diffusion_models/prx_pixel_t2i/src/__init__.py

@@ -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

extensions_built_in/diffusion_models/prx_pixel_t2i/src/pipeline.py

@@ -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]

extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py

@@ -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)

ui/src/app/jobs/new/options.ts

@@ -1004,6 +1004,27 @@ export const modelArchs: ModelArch[] = [
     ],
     hasMultiLinePrompts: true,
   },
+  {
+    name: 'prx_pixel',
+    label: 'PRXPixel (pixel space)',
+    group: 'image',
+    defaults: {
+      'config.process[0].model.name_or_path': ['Photoroom/prxpixel-t2i', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
+      'config.process[0].network.conv': [undefined, 16],
+      'config.process[0].network.conv_alpha': [undefined, 16],
+      'config.process[0].model.low_vram': [true, false],
+    },
+    disableSections: [
+      'network.conv',
+    ],
+    additionalSections: [
+      'model.low_vram',
+      'model.layer_offloading',
+    ],
+  },
 ].sort((a, b) => {
   // Sort by label, case-insensitive
   return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });