PRX Part 3: Training a Text-to-Image Model in 24 Hours
Hugging Face and Photoroom have demonstrated that a usable text-to-image model can be trained from scratch in just 24 hours. By combining several modern architectural and training optimizations, the team achieved this using 32 H200 GPUs with a total compute budget of approximately $1,500.
Pixel-Space Training and X-Prediction
The model utilizes the x-prediction formulation, which allows for training directly in pixel space, eliminating the need for a Variational Autoencoder (VAE). To keep sequence lengths computationally manageable, the team used a patch size of 32 and a 256-dimensional bottleneck in the initial token projection layer.
At 512px resolution, the sequence length is 256 tokens, and at 1024px, it increases to 1024 tokens. The training schedule began directly at 512px and concluded with fine-tuning at 1024px.
Perceptual Losses for Visual Quality
Predicting pixels directly enables the use of classical computer vision perceptual losses to improve convergence speed and visual quality. The team implemented two auxiliary losses on top of the standard flow matching objective:
- LPIPS: Used with a weight of 0.1 to capture low-level perceptual similarity.
- DINO-based perceptual loss: Using DINOv2 with a weight of 0.01 to provide a stronger semantic signal.
These losses were applied to pooled full images across all noise levels, rather than patch-wise features, to achieve better empirical results.
Efficient Token Routing with TREAD
To reduce the computational cost per step, the team employed TREAD (Token Routing for Efficient Architecture-agnostic Diffusion Training). This method randomly selects a fraction of tokens to bypass a contiguous chunk of transformer blocks before being re-injected.
Specifically, 50% of the tokens were routed from the second block to the penultimate block of the transformer. To mitigate potential quality degradation under vanilla Classifier-Free Guidance (CFG), a self-guidance scheme was implemented to guide using a dense vs. routed conditional prediction.
Representation Alignment and Optimizer
Representation alignment was achieved using REPA, with DINOv3 serving as the teacher model. The alignment loss was applied once at the 8th transformer block with a loss weight of 0.5. To maintain consistency with TREAD routing, the alignment loss was only computed on non-routed tokens.
For optimization, the team used the Muon optimizer for 2D parameters (matrices) and Adam for all non-2D parameters (biases, norms, embeddings).
| Optimizer Group | Target Parameters | Key Hyperparameters |
|---|---|---|
| Muon | 2D parameters | lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5 |
| Adam | Non-2D parameters | lr=1e-4, betas=(0.9, 0.95), eps=1e-8 |
Training Data and Schedule
Training was conducted on three synthetic datasets:
- Flux generated (1.7M images)
- FLUX-Reason-6M (6M images)
- midjourney-v6-llava (1M images), re-captioned using Gemini 1.5 for consistency.
The training schedule consisted of 100k steps at 512px with a batch size of 1024, followed by 20k steps at 1024px with a batch size of 512 (without REPA).
Results and Implications
While the final model exhibits some texture glitches and occasional anatomical errors, the team describes it as "clearly usable" with strong prompt following and a consistent aesthetic. The authors conclude that remaining issues are likely due to undertraining and limited data diversity rather than structural flaws in the recipe.
Photoroom has open-sourced the training code and experimental framework via the PRX GitHub repository.