MangaNinja Documentation

The ReferenceAttentionControl mechanism works by intercepting the forward pass of BasicTransformerBlock modules within the UNet.

1. Write Mode: As the UNet processes the reference image, the ReferenceAttentionControl hooks into the transformer blocks and appends the normalized hidden states to a bank attribute within each module.
2. Read Mode: When processing the target image, the hooks retrieve the stored features from the bank. It then performs a modified self-attention where the query comes from the target image, but the keys and values are augmented with the features from the reference bank.
3. Fusion Scopes:
   • "midup": Only modifies the mid_block and up_blocks. This is often more efficient and focuses on structural/color fusion at higher resolutions.
   • "full": Modifies every BasicTransformerBlock found in the entire UNet hierarchy.
4. Point Embeddings: The system supports point_bank_ref and point_bank_main, which allow for more precise, point-based reference following by filtering embedding matrices to match the specific spatial dimensions of the attention blocks.

When initializing Transformer2DModel, use the following parameters:

The forward method performs the denoising pass through the UNet. It supports various conditioning types including text, images, and additional residuals for ControlNet or T2I-Adapters.

Key Arguments:

• sample (torch.FloatTensor): Noisy input tensor of shape (batch, channel, height, width).
• timestep (torch.FloatTensor | float | int): The number of timesteps to denoise.
• encoder_hidden_states (torch.FloatTensor): Encoder hidden states of shape (batch, sequence_length, feature_dim).
• class_labels (torch.Tensor, optional): Class labels for conditioning.
• attention_mask (torch.Tensor, optional): Mask of shape (batch, key_tokens) (1 to keep, 0 to discard).
• encoder_attention_mask (torch.Tensor, optional): Cross-attention mask of shape (batch, sequence_length).
• added_cond_kwargs (dict, optional): Dictionary for additional embeddings. Depending on addition_embed_type in config, it may require:
  • image_embeds: For text_image, image, or image_hint modes.
  • text_embeds: For text_time mode.
  • time_ids: For text_time mode.
  • hint: For image_hint mode.
• down_block_additional_residuals (tuple[torch.Tensor], optional): Residuals for ControlNet.
• mid_block_additional_residual (torch.Tensor, optional): Residual for the middle block.
• down_intrablock_additional_residuals (tuple[torch.Tensor], optional): Residuals for T2I-Adapters.
• return_dict (bool, optional, defaults to True): If True, returns UNet2DConditionOutput; otherwise returns a tuple containing the sample.

# Example: Basic forward pass
output = model(
    sample=noisy_latents,
    timestep=50.0,
    encoder_hidden_states=text_embeddings
)

# If return_dict=False, access the sample directly
output_tuple = model(
    sample=noisy_latents,
    timestep=50.0,
    encoder_hidden_states=text_embeddings,
    return_dict=False
)
sample = output_tuple[0]

The Transformer2DModel is a 2D Transformer designed for image-like data. It can process three types of inputs based on the configuration provided:

1. Continuous Input: Standard images with shape (batch_size, num_channels, height, width). Requires setting in_channels and ensuring patch_size is None.
2. Vectorized Input: Quantized image embeddings with shape (batch_size, num_vector_embeds). Requires setting num_vector_embeds and ensuring patch_size is None.
3. Patch Input: Requires setting both in_channels and patch_size.

Note: You cannot define both in_channels and num_vector_embeds simultaneously, nor can you define both num_vector_embeds and patch_size.

from src.models.transformer_2d import Transformer2DModel

# Example for continuous input
model = Transformer2DModel(
    num_attention_heads=16,
    attention_head_dim=88,
    in_channels=4,
    num_layers=12,
    cross_attention_dim=768
)