OmniGen Next Generation Image Generation on Cloud GPUs

• Python: The content of this article is highly technical. We recommend this piece to readers experienced with both Python and basic concepts in Deep Learning. For new users, this article may be a good place to start.
• Cloud GPU: Running FLUX.1 will require a sufficiently powerful GPU. We recommend at least 40 GB VRAM machines at the minimum.

OmniGen is comprised of a two part composite of a Variational AutoEncoder (VAE) and large pretrained Transformer. The VAE functionally extracts visual features from the images continuously, while the Transformer generates images based on the input conditions. Specifically, these are the VAE from Stable Diffusion XL which was frozen during training, and the Transformer model was initialized with Microsoft’s Phi-3. This allows a functional connection between the strength of the VAE with the a Transformer that has inherited a signficant understanding of textual processing capability. Put together, this creates a simple but strong pipeline that prevents the need for additional encoders in the model, which simplifies the pipeline significantly. OmniGen inherently encodes conditional information by itself. “Furthermore, OmniGen jointly models text and images within a single model, rather than independently modeling different input conditions with separate encoders as in existing works which lacks interaction between different modality conditions” Source.

For the attention mechanism, they use a modified version of the causal attention mechanism. Specifically, the mechanism works by applying causal attention to each element in the sequence, and simultaneously applu bidirectional attention within each image sequence. This makes it possible for each patch to “pay attention” to the other parts of the image and ensure that each image can only consider image or text sequences that have previously appeared Source.

To generate an image, they randomly sample a Gaussian noise and apply flow matching to predict the target velocity, and iterate the set number of inference steps to generate the images latent representation. The VAE then decodes the value into the final image output Source.

OmniGen is capable of numerous tasks, but, more importantly, abstracts away extra steps from the increasibly long process of image generation/editing with AI technologies. Let’s briefly overview each of them before jumping into the coding demo.

• Text-to-image generation: Like Stable Diffusion or FLUX, OmniGen is perfectly capable of generating high quality images on its own with a high degree of performance. In our experience, the quality is very familiar to that of using the baseline Stable Diffusion XL. It will be interesting to see if the same process could be applied to the more modern FLUX and SD 3.5 Large models to give this same capability to those models

• Text based image editing: OmniGen makes it easy to edit images in a single step with text instruction. It uses an LLM style prompting format, so asking the model to change the image subject’s hair color is a simple request

• Image compositing: OmniGen makes it easy to combine two subject matters seamlessly into novel environments

• Depth/pose estimation: OmniGen comes integrated with the same technologies that make ControlNet’s so effective, and are capable of doing the extraction of the pose themselves

• and much more! OmniGen is probably the most versatile single model pipeline ever released, be sure to try all the examples provided by the demo

OmniGen is a fascinating step forward for image generation models. In particular, we are impressed by the consolidation of the entire pipeline into a single model file capable of such a diverse array of tasks, including editing, image composition, and much more. We look forward to the release of the next versions of OmniGen in the coming months as the pipeline framework is spread to other models.