transfusion-pytorch

The model accepts a list of sequences. Each sequence is a list containing interleaved text and modality tensors.

Data Types:

• Text: Represented by torch.long tensors.
• Modalities: Represented by torch.float tensors.

Multi-modality indexing: When using multiple modalities, you can pass a tuple (modality_index, Tensor) for float tensors to specify which modality the tensor belongs to. The modality_index corresponds to the index in the dim_latent tuple provided during initialization.

Install the package using pip:

$ pip install transfusion-pytorch

To run the included training examples, install with the [examples] extra:

$ pip install .[examples]

If you encounter issues with safetensors, ensure the following dependencies are installed and up to date:

$ pip install -U diffusers transformers accelerate scipy ftfy safetensors

To train the model, pass the formatted input list to the model instance to compute the loss, then perform backpropagation.

To generate new content after training:

• Use model.sample() to generate a multi-modal sample.
• Use model.generate_text_only(text_tokens, length) for text-only generation (useful for language-first pretraining).
• Use print_modality_sample(sample) to inspect/print modality samples.

When return_loss=True is passed to the model, it computes a composite loss function:

1. Text Loss: Autoregressive cross-entropy loss on text tokens (using ignore_index for non-text positions).
2. Flow Loss: MSE loss between the predicted flow and the target flow (calculated from the noised modality).
3. Velocity Consistency Loss (Optional): If need_velocity_matching is active, it penalizes differences between the current flow and an EMA model's flow to 'straighten' the flow trajectories.
4. Reconstruction Loss (Optional): If has_recon_loss is enabled, it calculates MSE between the reconstructed modality and the original target.

The final loss is a weighted sum of these components, controlled by text_loss_weight, flow_loss_weight, velocity_consistency_loss_weight, and reconstruction_loss_weight.

In Transfusion, data is represented as ModalitySamples. A sample is a list containing text tokens (as torch.long tensors) and modalities (as torch.float tensors or tuples).

When using tuples for modalities, the format is (modality_type, tensor).

Key type definitions:

• ModalitySample: list[Int[''] | Int['_'] | Float['...'] | tuple[int, Float['...']]]
• ModalityTokenTransform: str | Callable | None
• RawModalityPositions: list[list[tuple[int, int, int]]] (representing [type, offset, length] for each modality).

The model uses a joint transformer and flow-matching system. It intersperses text tokens with modality tokens in a single sequence.

The sequence structure for a modality includes:

1. Meta Information: A meta_id followed by a character-tokenized string representing the modality's shape (e.g., '3,224,224').
2. Start of Modality (SOM): A special som_id token.
3. Modality Content: The actual latent tokens (often flattened or packed).
4. End of Modality (EOM): A special eom_id token.

Positional Embeddings:

• Text: Uses standard rotary embeddings derived from the sequence position.
• Modalities: Can use Axial Positional Embeddings if add_pos_emb is enabled. These are factorized embeddings that account for the multi-dimensional structure (like height/width) of the modality latents.

The Transfusion model handles two primary types of data in its input sequences:

1. Text: Represented as torch.long tensors. These are treated as discrete tokens.
2. Modalities (e.g., images, audio): Represented as torch.float tensors.

When passing modalities, you can provide them as raw tensors or as a tuple (modality_type, modality_tensor). If a torch.float tensor is provided without a tuple, it is implicitly treated as modality type 0.

Key constraints:

• Text tensors must be 1D (torch.long).
• Modality tensors must match the dimensions and latent channel size specified in the model's modality configuration.

Transfusion treats text and continuous modalities (like images) as part of a single sequence.

1. Text: Represented by discrete tokens. The model uses standard causal transformer attention.
2. Modalities: Represented by continuous latents. The model uses Flow Matching to generate these latents.
3. Transitions: The model uses special tokens to signal the start and end of modalities:
   • [som] (Start of Modality): Triggers the transition from text decoding to continuous flow sampling.
   • [eom] (End of Modality): Signals the end of the continuous sequence, returning the model to text decoding.
4. Shape Prediction: Before generating a modality, the model predicts a 'meta string' (via a character tokenizer) that describes the target shape (e.g., "256,256"). This shape is then used to initialize the noise for the flow matching process.

To inspect or manipulate the specific components of a particular modality (e.g., its encoder, decoder, or dimension settings), use get_modality_info(modality_type). This returns a ModalityInfo object containing:

• encoder: The Module used to encode the modality.
• decoder: The Module used to decode the modality.
• latent_to_model: Projection from latent space to transformer dimension.
• model_to_latent: Projection from transformer dimension to latent space.
• add_pos_emb: Boolean indicating if axial positional embeddings are used.
• pos_emb_mlp: The MLP used to generate axial positional embeddings.
• num_dim: The number of dimensions for the modality.
• dim_latent: The dimensionality of the latent space.
• default_shape: The fallback shape for this modality.
• som_id / eom_id: The start and end tokens for this modality type.

The transfusion_pytorch package provides the core Transfusion model class along with utilities for data loading and inspecting modality samples. You can import these directly from the package root.

from transfusion_pytorch import (
    Transfusion,
    print_modality_sample,
    create_dataloader
)