Welcome aMUSEd: Efficient Text-to-Image Generation

We’re excited to present an efficient non-diffusion text-to-image model named aMUSEd. It’s called so because it’s a open reproduction of Google's MUSE. aMUSEd’s generation quality is not the best and we’re releasing a research preview with a permissive license.

In contrast to the commonly used latent diffusion approach (Rombach et al. (2022)), aMUSEd employs a Masked Image Model (MIM) methodology. This not only requires fewer inference steps, as noted by Chang et al. (2023), but also enhances the model's interpretability.

Just as MUSE, aMUSEd demonstrates an exceptional ability for style transfer using a single image, a feature explored in depth by Sohn et al. (2023). This aspect could potentially open new avenues in personalized and style-specific image generation.

In this blog post, we will give you some internals of aMUSEd, show how you can use it for different tasks, including text-to-image, and show how to fine-tune it. Along the way, we will provide all the important resources related to aMUSEd, including its training code. Let’s get started 🚀

Table of contents

• How does it work?
• Using in diffusers
• Fine-tuning aMUSEd
• Limitations
• Resources

We have built a demo for readers to play with aMUSEd. You can try it out in this Space or in the playground embedded below:

How does it work?

aMUSEd is based on Masked Image Modeling. It makes for a compelling use case for the community to explore components that are known to work in language modeling in the context of image generation.

The figure below presents a pictorial overview of how aMUSEd works.

During training:

• input images are tokenized using a VQGAN to obtain image tokens
• the image tokens are then masked according to a cosine masking schedule.
• the masked tokens (conditioned on the prompt embeddings computed using a CLIP-L/14 text encoder are passed to a U-ViT model that predicts the masked patches

During inference:

To learn more about aMUSEd, we recommend reading the technical report here.