Add support for PRX Pixel T2I

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

View File

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