Add support for PRX Pixel T2I
This commit is contained in:
53
extensions_built_in/diffusion_models/prx_pixel_t2i/README.md
Normal file
53
extensions_built_in/diffusion_models/prx_pixel_t2i/README.md
Normal 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.
|
||||
Reference in New Issue
Block a user