ai-toolkit/extensions_built_in/diffusion_models/prx_pixel_t2i

PRXPixel (Photoroom PRX-7B, pixel-space text-to-image)

ai-toolkit integration for 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: the transformer is vendored in src/transformer_prx.py and a minimal preview sampler lives in 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:

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

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 for the generic lifecycle, registration and LoRA conventions.