Welcome back to the PRX series! In Parts 1-3, we dove into model architectures, training design, and a thrilling 24-hour speedrun. Now, it’s time to reveal the unsung hero that quietly powers it all: our data strategy. While perhaps not the flashiest component, building a robust data pipeline was absolutely crucial to PRX’s remarkable quality.
It was a journey of careful planning, unexpected challenges, and invaluable lessons learned the hard way. This installment pulls back the curtain on our approach, sharing what worked, what we’d refine next time, and the key insights that emerged along the way. In a nutshell, we assemble training data from a blend of public and internal datasets, meticulously re-caption images using a Vision-Language Model (VLM), and transform this into the streamable corpus PRX was trained on.
The Philosophy Behind Our Data Strategy
Our guiding principle was simple yet profound: build a large, diverse dataset for pre-training. At this foundational stage, the model is essentially learning “how the world looks,” absorbing visual concepts, objects, scenes, composition, and lighting. The goal here is broad coverage and diversity, not per-image perfection; a wide, representative corpus teaches the model far more than a smaller, “prettier” one.
In fact, over-filtering for aesthetics at this stage can be detrimental, narrowing the distribution and depriving the model of essential concepts. Polishing generations comes later, during fine-tuning with small, meticulously curated sets. For pre-training, breadth is paramount; for fine-tuning, taste dictates.
We combine public and internal datasets, prioritizing existing curation efforts rather than reinventing the wheel. Our pragmatic approach allowed us to quickly assemble a robust starting point for pre-training a 7B model. While perhaps not the absolute best theoretical dataset, it was solid, lightweight, and incredibly effective.
Crucially, our captioning philosophy emphasizes long, highly accurate captions that describe everything within an image. We observed in Part 2 that switching from short to long captions dramatically improved sample quality. This meticulous captioning transforms what might be considered “noise”—like screenshots or logos—into controllable attributes that the model can learn to either generate or omit based on prompts.
Data Formats and Efficiency: Lance vs. MDS
For distributed training, we’ve long relied on Mosaic Streaming and Mosaic Data Shards (MDS) due to their low-maintenance, flexible, and high-performance nature. MDS datasets excel at mixing, shuffling, and direct streaming from object storage like S3 or GCS. However, their rigidity for modifications led us to adopt another powerful tool: Lance.
Lance is a columnar data format optimized for feature engineering and dataset curation, offering cheap predicate pushdown, scalar indexes, and vector search. It’s the ideal tool for building and exploring datasets with billions of rows. This synergy is key: Lance for building and curating, MDS for efficient streaming during training.
Previously, we pre-computed and stored text latents, but with the switch to Qwen3-VL as our text encoder, we now compute them on the fly. While this adds a small 3-4% throughput cost, it significantly reduces MDS shard sizes, allowing us to store the entire pre-training dataset on SSD-backed shared filesystems. This also grants us the flexibility to swap text encoders without re-writing terabytes of data, a massive advantage for iterative development.
Image Encoding and Perceptual Quality
We opted to encode all images as JPEG at quality 92, rather than a lossless format like PNG. This decision wasn’t arbitrary; we rigorously measured its impact. Given that most real-world images are already JPEG-compressed, the critical question was whether one more re-encode would cause perceptible degradation.
Our tests showed that at quality 92, the first re-encode is essentially imperceptible, and subsequent cycles add almost no further degradation as JPEG quickly converges. Even after 10 cycles, images remained well within the imperceptible range, while PNGs would be 3–10× larger for no perceptual gain. For example, averaging over 100 images at quality 92, we found PSNR values around 38.60 dB and LPIPS around 0.007, indicating minimal quality loss.
Most importantly, we confirmed that training on JPEGs does not negatively impact model outputs. We trained two identical PRX models at 1024px, one on PNGs and one on JPEGs at quality 92. Both performed equally well, producing generations that were practically indistinguishable. Our analysis revealed that only about one in ten generations from either model showed any detectable quantization structure, and the differences were negligible. While we still use PNGs for artifact-sensitive data in other Photoroom models (like our custom AI Shadows), high-quality JPEG storage is more than sufficient for large-scale text-to-image training.
Building and Exploring the Dataset in Lance
Interactively exploring Parquet tables with billions of rows is challenging, which is precisely why Lance became our go-to. It enables indexing, querying, and browsing at scale. We leverage Ray Data for ingest, reading source tables in parallel and writing numerous Lance table fragments across the cluster.
A crucial lesson learned concerned fragmentation: too many small fragments can severely degrade query performance. We initially targeted 100,000 rows per fragment, leading to thousands of tiny fragments and sluggish queries. After compacting to approximately a million rows per fragment, bringing us to around a thousand fragments for our hundreds of millions of rows, query speeds dramatically improved. While Lance supports distributed compaction, it’s wise to avoid over-fragmentation from the outset.
Although we re-caption every image for consistency, existing captions and embeddings are invaluable for initial exploration. Pre-existing metadata, such as CLIP embeddings, allows for immediate full-text and nearest-neighbor searches in Lance. This capability lets us quickly navigate the data, assess its quality, and determine necessary filtering even before running our own captioning pipeline.
With data queryable, we rapidly profiled it. For instance, the resolution distribution helped us set cutoffs; we discard anything below an effective cutoff of 384² pixels (approx. 147k pixels) with an aspect ratio between 0.5 and 2.0. We built a small UI leveraging Lance’s full-text and nearest-neighbor search capabilities over indexed caption and embedding columns. This enabled real-time textual search and visual navigation through image clusters.
This interactive browsing was instrumental in qualitatively assessing diversity and quality, allowing us to spot issues early. We uncovered uninformative baseline captions, a surprising amount of non-photographic content (screenshots, documents), and near-duplicate images. These observations directly influenced the rest of our pipeline, reinforcing the need for comprehensive re-captioning with longer, accurate descriptions.
Source: Hugging Face Blog
