# transfusion-pytorch: one transformer for text tokens and flow-matched modalities

> A PyTorch implementation of MetaAI's Transfusion that swaps diffusion for flow matching and extends the idea to any number of modalities. It is a research library, not a pipeline, and the README tells you to bring your own encoder and decoder.

**lucidrains/transfusion-pytorch** — Pytorch implementation of Transfusion, "Predict the Next Token and Diffuse Images with One Multi-Modal Model", from MetaAI

- Repository: https://github.com/lucidrains/transfusion-pytorch
- Stars: 1,404 · Forks: 75
- Language: Python
- License: MIT
- Published: 2026-09-10 · Updated: 2026-09-10 · Language: en
- Canonical page: https://hysenlabs.com/projects/lucidrains-transfusion-pytorch

## What transfusion-pytorch solves, and who it is actually for

Most multimodal systems bolt a diffusion model onto a language model and pass latents between them. Transfusion, the MetaAI paper this repository implements, argues for one transformer that predicts text tokens with next-token prediction and generates continuous modalities with a diffusion-style objective, in the same sequence. lucidrains/transfusion-pytorch is the PyTorch port of that idea. The README states one deliberate deviation: diffusion is replaced with flow matching, citing the success of Flux from Black Forest Labs, while the paper's title is kept because "Transflow does not have the same ring".

The target reader is someone who wants to experiment with that architecture rather than consume a product. The package is a model class plus training scripts. There is no inference server, no pretrained weights, no dataset downloader. The repository root holds train_image_only.py, train_mnist.py, train_latent_with_text.py, train_self_flow.py and several more, each a self-contained experiment. If you want a working image generator today, this is the wrong dependency. If you want to test whether a single transformer can hold text and continuous latents in one context, the surface is small enough to read in an afternoon.

## How the mechanism works: one sequence, two objectives

The central design choice is that the data type of a tensor decides how it is handled. The README is explicit: "any torch.long is text, torch.float is modalities". A training example is a list of tensors, alternating between integer token sequences and float latent tensors. The model computes a loss over the whole list, and loss.backward() works as usual.

When you pass multiple modalities, dim_latent becomes a tuple, one entry per modality, and modality_default_shape gives a fallback shape for each in case the language model does not produce a valid modality shape. A float tensor is then wrapped as a tuple whose first element is the modality index, for example (1, randn(6, 192)) for the second modality.

Encoding is pluggable. In the third README example, modality_encoder and modality_decoder are set to nn.Conv2d layers, and channel_first_latent = True tells the model the latent layout. That example is a mock: the README uses nn.Conv2d(3, 384, 3, padding = 1) as a stand-in, which is a strong hint about where the real work sits. The transformer itself is configured through a nested dict with dim and depth. There is no built-in VAE.

## Installing transfusion-pytorch and running a first forward pass

The README gives a single install line. The package requires Python 3.10 or newer according to pyproject.toml, and pulls in torch>=2.0, einops, einx, jaxtyping, beartype, torchdiffeq, rotary_embedding_torch, axial-positional-embedding, ema-pytorch, torch-einops-utils, loguru and tqdm.

```bash
pip install transfusion-pytorch
```

A minimal single-modality model takes a token vocabulary size, a latent dimension, a fallback modality shape and a transformer config. The README's first example is:

```python
from torch import randint, randn
from transfusion_pytorch import Transfusion

model = Transfusion(
    num_text_tokens = 256,
    dim_latent = 384,
    modality_default_shape = (4,),
    transformer = dict(
        dim = 512,
        depth = 8
    )
)
```

Training data is a list of lists. Each inner list mixes torch.long token tensors and torch.float latent tensors. Calling the model returns a loss, and the README shows loss.backward() immediately after. After training, model.sample() returns one multimodal sample.

To run the scripts in the repository root, the README says to install the examples extra first, and adds a fallback command if safetensors misbehaves:

```bash
pip install .[examples]
pip install -U diffusers transformers accelerate scipy ftfy safetensors
```

The examples extra adds adam-atan2-pytorch, datasets and diffusers. Note that pip install .[examples] requires a local checkout; the PyPI package alone does not ship the training scripts.

## Batched sampling with sample_many, and forcing a modality first

Sample generation has two entry points. sample_one walks a single sample through a state machine serially. sample_many is the batched equivalent: the README states that all samples currently decoding text share a single kv-cached forward pass, and all samples currently decoding a modality share a single joint odeint trajectory, with one forward pass per ODE evaluation while each sample keeps its own shape, length and modality type. The kv cache is always used on that path, and each sample carries its own max length budget, so an early finisher does not block the batch.

The batch is a list of prompts. Each prompt can be a raw text tensor, a (modality_type, tensor) tuple, an interleaved list, or None for an empty prompt. The call returns a list of ModalitySample objects, one per prompt:

```python
prompts = [
    torch.randint(0, 256, (16,)),
    (0, torch.randn(4, 384)),
    None,
    [torch.randint(0, 256, (8,)), (0, torch.randn(6, 384))],
]

samples = model.sample_many(prompts, max_length = 2048, cfg_scale = 3.)
```

Classifier-free guidance is available through cfg_scale; the README credits Pranoy for adding it. There is also force_modality_at_start, which skips text decoding at the start and makes generation begin with a chosen modality, optionally with a forced shape. The README uses this for an in-context human-to-robot case: a human demonstrates a task on video, and robot actions are generated in-context without finetuning. Whether that works at your scale is an empirical question the README does not answer.

## Where transfusion-pytorch will disappoint you

The biggest gap is the one the README's own example exposes. modality_encoder and modality_decoder are nn.Conv2d layers in the demonstration, which means the library assumes you already have something that turns pixels into latents and back. There is no VAE, no tokenizer, no pretrained checkpoint and no dataset loader in the package. The repository layout confirms this: data/ and a set of train_*.py scripts, nothing that downloads or prepares a corpus.

Second, the project is a beta by its own classification. pyproject.toml carries the classifier "Development Status :: 4 - Beta", and the version in pyproject.toml is 0.20.0 while the most recent release listed is 0.16.3. That gap is worth checking before you pin a version, because the README documents behaviour such as sample_many that may not exist in the version pip resolves.

Third, there is no documented rollback path, no migration note and no changelog. The README does not document breaking changes between 0.16.x and the current source. If you build on this, pin the exact version and keep your own notes on which API surface you depend on. Finally, sample_many is described as sharing one ODE trajectory across the batch. That is an efficiency claim about batching, not a claim about sample quality, and the README offers no evaluation numbers either way.

## The alternative: keep the diffusion model separate

The obvious alternative is the two-model arrangement that Transfusion was written against: a language model that emits text and a separate diffusion model conditioned on that text, communicating through embeddings or latents. The difference is architectural, not cosmetic. In the two-model design each component has its own training loop, its own checkpoint and its own failure modes, and you can swap the image model without retraining the language model. In transfusion-pytorch the two objectives share one transformer and one loss, so a change to the modality pathway touches the text pathway too.

That coupling is the point of the paper. It is also why the repository is a research artifact rather than a production dependency. If your goal is to ship an image feature, a separate diffusion model with published weights will get you there faster. If your goal is to test whether interleaved prediction and generation in one context helps on your data, the single-model design is the thing you are trying to measure, and the two-model baseline is what you compare against.

## Licence, maintenance and upgrade cost

The licence is MIT, declared both in the LICENSE file and in the pyproject.toml classifier "License :: OSI Approved :: MIT License". MIT is permissive: it allows commercial use and modification provided the copyright notice and permission notice are retained. This is a factual description of the licence text, not legal advice; if your organisation has a policy on copyleft or on model weights, route it through whoever handles that.

The repository is not archived, and the last push was on 2026-09-04. The most recent release listed is 0.16.3 from 2026-01-27, with 0.16.2 and 0.16.1 earlier that month. The version string in pyproject.toml is 0.20.0, so releases and source have diverged. The dependency list is long and includes fast-moving packages (einx>=0.4.0, rotary_embedding_torch>=0.8.4, torch-einops-utils>=0.1.22), which means an upgrade can pull in changes from several transitive dependencies at once. Budget for pinning the full lockfile rather than just transfusion-pytorch.

## Conclusion

Adopt transfusion-pytorch if you already have a VAE or encoder that produces latents and you want to prototype interleaved text and modality generation in a single transformer, including the in-context video-to-action example the README sketches. Do not adopt it if you need an end-to-end image pipeline: the README's own example wires in nn.Conv2d as a mock encoder and decoder, and there is no checkpoint, tokenizer or dataset loader in the package. Before committing, verify that the version on PyPI matches the 0.20.0 in pyproject.toml, since the most recent release listed is 0.16.3.

## FAQ

### How do I install transfusion-pytorch?

Install it from PyPI with pip install transfusion-pytorch. It requires Python 3.10 or newer and pulls in torch>=2.0 along with einops, einx, jaxtyping and several other dependencies. To run the training scripts in the repository root you also need the examples extra, installed from a local checkout.

### Does transfusion-pytorch include an image encoder or a pretrained model?

No. The README's encoding example uses nn.Conv2d layers as a mock encoder and decoder, passed in through modality_encoder and modality_decoder, which means you supply the real VAE or encoder yourself. There is no pretrained checkpoint, tokenizer or dataset loader in the package.

### What is the difference between sample_one and sample_many in transfusion-pytorch?

sample_one walks a single sample through the generation state machine serially. sample_many is the batched equivalent: samples currently decoding text share one kv-cached forward pass, and samples decoding a modality share one joint odeint trajectory, with each sample keeping its own shape, length and modality type.

## Sources

- [Issues](https://github.com/lucidrains/transfusion-pytorch/issues)
- [License: MIT](https://github.com/lucidrains/transfusion-pytorch/blob/main/LICENSE)
- [lucidrains/transfusion-pytorch on GitHub](https://github.com/lucidrains/transfusion-pytorch)
- [README](https://github.com/lucidrains/transfusion-pytorch/blob/main/README.md)
- [Releases](https://github.com/lucidrains/transfusion-pytorch/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/lucidrains-transfusion-pytorch