Transfusion: One Model to Predict Text and Generate Images

You've seen models that handle text. You've seen models that handle images. What if one model could do both—without needing separate modules or complex routing?

That's the idea behind Transfusion, a new architecture that combines text prediction (like a language model) with image diffusion (like a stable diffusion model), all inside a single transformer. No separate encoder for text, no dedicated decoder for images. Just one unified model that learns to both predict the next token and denoise pixels.

And now there's a clean, well-documented PyTorch implementation you can actually run.

What It Does

Transfusion is a transformer-based model trained on multimodal sequences (text + images). It processes tokens and image patches in the same sequence, and uses two different loss functions:

• Causal language modeling loss for text tokens (predict the next token)
• Diffusion loss for image patches (predict the noise added to patch embeddings)

During training, the model sees examples like "A cat sitting on a [image patches]" and learns to both predict the next word and reconstruct the image patches. During inference, you can feed it text and generate an image, or feed it an image and generate text.

The GitHub repo (lucidrains/transfusion-pytorch) provides a minimal but functional implementation with clear code, attention masking, and support for different image sizes.

Why It's Cool

A few things stand out:

No modality-specific encoders. Many multimodal models (like CLIP, LLaVA, or DeepFloyd) rely on separate encoders for images or text, then fuse them. Transfusion does everything in a single transformer, which is both simpler and potentially more efficient.

You get two modes from one model. The same weights can predict text and generate images. That's not just convenient—it forces the model to learn shared representations between modalities.

Scales naturally. Because it uses the same attention mechanism and positional encoding for both text tokens and image patches, you can increase the number of image patches or sequence length without major architectural changes.

Practical for small experiments. The code is well-structured and easy to modify. You can train a tiny version on a single GPU to test ideas, then scale up.

How to Try It

The repo is straightforward. Here's the quick start:

pip install transfusion-pytorch

Then, in Python:

import torch
from transfusion_pytorch import Transfusion

# Create a small model (adjust sizes for your GPU)
model = Transfusion(
    num_text_tokens=256,   # Example: character-level vocab
    dim_latent=512,
    dim=512,
    depth=6,
    dim_head=64,
    heads=8,
    ff_mult=4,
    image_size=64,
    patch_size=16,
    channels=3,
    alibi_pos_bias=False,
    ff_glu=True
)

# Dummy data: text tokens + image patches
text_tokens = torch.randint(0, 256, (2, 256))
image_patches = torch.randn(2, 3, 4, 4)  # (batch, channels, grid_h, grid_w)
image_mask = torch.ones(2, 16).bool()    # 4x4 = 16 patches

loss = model(text_tokens, image_patches, image_mask=image_mask)
loss.backward()

For actual training, you'll need a dataset of text-image pairs, but the repo gives you the core loop.

If you want to see it work without training, check the README—there are examples for sampling (text→image and image→text).

Final Thoughts

Transfusion is one of those ideas that feels obvious in hindsight. Why not use the same architecture for both modalities? The implementation here is minimal but solid—you can read the source in a few minutes and understand exactly what's happening.

If you're interested in multimodal models, or just want to see how far you can push a single transformer, give it a try. It's a great starting point for experiments, and the code is clean enough to hack on.

Give it a star if you find it useful, and let us know what you build.

Found on @githubprojects