transfusion-pytorch: one transformer that predicts tokens and flow-matches latents

The problem: text and images do not share a prediction objective

Autoregressive language models predict a categorical distribution over the next token. Image generators predict a continuous distribution over pixels or latents. Most multi-modal systems bolt these together: a text model calls out to a diffusion model, or a diffusion model is conditioned on text embeddings from a frozen encoder. Transfusion, the MetaAI paper this repository implements, argues for one transformer trained on a single sequence where text tokens and image latents are interleaved, with the loss switching depending on what the next element is. The README states the repository substitutes flow matching for diffusion, citing the success of Flux from Black Forest Labs, while keeping the original paper's title. The target user is a researcher or research engineer who wants to experiment with that unified objective rather than assemble a pipeline of separate models. This is not a library for adding image generation to an existing product. It is a model definition plus training loop scaffolding, and the README treats it that way: the examples are train_{example_name}.py scripts in the project root, not a serving stack.

How the single-transformer mechanism distinguishes text from latents

The dispatch rule is by dtype. The README repeats it as a comment in two examples: any torch.long is text, torch.float is modalities. A training batch is a list of sequences, and each sequence is a list whose elements are either an integer tensor of token ids or a float tensor of latent values. The model computes a loss over the whole interleaved sequence, and loss.backward() is all the caller does. For multiple modalities, a float tensor can be wrapped as a tuple[int, Tensor] where the first position is the modality index, and dim_latent becomes a tuple of per-modality latent dimensions with modality_default_shape giving a fallback shape per modality. The fallback matters because the language model may not produce a valid modality shape, per the README's comment on modality_default_shape. That is the whole interface: dtype tells the model what kind of prediction to make, and the modality index tells it which latent space to flow-match in. Sampling is a state machine. sample_one walks a single sample through it serially, alternating between token decoding and an odeint trajectory over the latent. sample_many is the batched equivalent: the README says 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 in sample_many, corresponding to the cache_kv = True path of sample_one.

Encoding and decoding modalities without writing a pipeline

If you already have an encoder and decoder, you pass them in. The README's example uses mock_encoder = nn.Conv2d(3, 384, 3, padding = 1) and a mirrored decoder, sets channel_first_latent = True and modality_default_shape = (4, 4), and then feeds raw image tensors such as randn(3, 8, 8) directly into the interleaved sequence. The model handles encoding into latents and decoding back out. print_modality_sample renders a sample. The channel_first_latent flag is per modality and can be a tuple when you have several, as the video-and-action example shows with channel_first_latent = (True, False). The design assumes you bring your own autoencoder. Nothing in the README suggests a bundled VAE or a pretrained image tokenizer, so the quality of your latents is your problem, not the library's. That is a reasonable boundary for a research implementation, and it is worth being explicit about: this repository gives you the sequence model and the training objective, not the representation learning.

Install and the shortest path to a training step

Installation is one command: pip install transfusion-pytorch. For the example scripts, the README says to run pip install .[examples] from the project root, and adds a fallback if safetensors misbehaves: pip install -U diffusers transformers accelerate scipy ftfy safetensors. A minimal model is Transfusion(num_text_tokens = 256, dim_latent = 384, modality_default_shape = (4,), transformer = dict(dim = 512, depth = 8)). The transformer sub-dict holds the backbone configuration. Training data is the nested list described earlier, and the call is loss = model(text_and_images) followed by loss.backward(). Text-only pretraining uses the same class: pass a tensor of shape (batch, seq) of token ids, and the README shows model.generate_text_only(text[:, :1], 1024) for sampling afterwards. Classifier-free guidance is available at sampling time via cfg_scale, shown as cfg_scale = 3. in the sample_many call. The README credits Pranoy for adding classifier free guidance. Forcing the first generated element to be a modality is done with force_modality_at_start on sample or sample_many, either as a modality type or as a (modality_type, shape) tuple. The README's example uses that for a human-to-robot setting where a video demonstration is followed by generated action chunks, described as in-context with no finetuning.

Where the abstraction leaks

The dtype dispatch is clever and it is also the sharpest edge in the API. A float tensor is always a modality, so any continuous quantity you want to treat as a regression target becomes a flow-matching target, and any integer you want to treat as a latent becomes a token. You cannot have a continuous scalar that is not a modality without wrapping it as one. The second edge is shape inference. modality_default_shape exists precisely because the model may not emit a usable shape, which means malformed shapes are an expected runtime condition rather than an error you can rule out by construction. The README does not describe what happens when a sample's generated shape disagrees with its default, only that the default is the fallback. Third, the repository is a model definition without a documented training recipe. There is no stated dataset, no hyperparameter table, no convergence guidance in the supplied material. The examples directory is referenced by filename pattern only. Fourth, sample_many shares one joint odeint trajectory across samples decoding a modality at the same time. That is a throughput decision, and the README does not state whether the shared trajectory is numerically identical to running each sample through sample_one. If you care about reproducing single-sample results at batch scale, that is the first thing to check. Finally, no pretrained weights are mentioned anywhere in the README. You are training from scratch.

How it differs from a diffusion-plus-LLM pipeline

The obvious alternative is the pipeline approach: a language model for text and a separate diffusion model conditioned on its embeddings, connected by an adapter. That architecture is well understood, each half can be trained and replaced independently, and pretrained components exist for both. The difference in approach is where the coupling lives. In a pipeline, the interface is a fixed embedding vector passed between two models that never share parameters or attention. In transfusion-pytorch, text tokens and modality latents occupy one sequence inside one transformer, and the objective switches per element by dtype. The practical consequences follow from that. You cannot swap out the image half without retraining the whole model, because there is no image half. You also do not need an adapter to align two representation spaces, because there is only one. Whether the unified model actually produces better cross-modal coherence is an empirical question the README does not answer, and this review cannot either. If your goal is to ship image generation this quarter, the pipeline route has pretrained checkpoints and this does not.

Maintenance, versioning and licence

The repository is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are preserved. That is a statement about the licence text, not legal advice; if you are shipping a product, have counsel read it. The project is not archived and the last push recorded is 2026-09-04. The release history shows 0.16.1, 0.16.2 and 0.16.3 within about a week of each other in January 2026, which suggests active iteration on a fast patch cadence rather than a stable API. For a research dependency that is fine. For anything pinned in a production environment it means you should pin an exact version and read the diff between patches, because the minor version number is not signalling interface stability. Upgrade cost is mostly the cost of re-reading the README and the example scripts after each bump. There is no changelog content in the supplied material, so the release notes themselves are the only record of what changed. The dependency surface is the usual PyTorch research stack plus, for examples, diffusers, transformers, accelerate, scipy, ftfy and safetensors, all of which move independently of this project.

Editorial conclusion

Adopt transfusion-pytorch if you are doing research on unified multi-modal generation and want a readable reference implementation of the Transfusion idea with flow matching, multi-modality support, and batched sampling already wired up. Do not adopt it if you need a documented, stable API for production inference, if you depend on pretrained checkpoints, or if you cannot afford to read the source to understand sampling behaviour. Before committing, verify three things: that sample_many's shared-trajectory batching matches the quality you get from sample_one, that force_modality_at_start with a (modality_type, shape) tuple behaves as the README describes for your modality layout, and that the encoder and decoder you plug in via modality_encoder and modality_decoder produce latents of the shape your modality_default_shape expects. The MIT licence lets you do what you want with the code; it does not give you pretrained weights, and the repository does not claim to ship any.

Official sources

1. Issues
2. License: MIT
3. lucidrains/transfusion-pytorch on GitHub
4. README
5. Releases