Skip to content

[docs] refactoring docstrings in ./src/diffusers/models/transformers/auraflow_transformer_2d.py#9715

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ahnjj wants to merge 3 commits into
base:main
Choose a base branch
from
Open

[docs] refactoring docstrings in ./src/diffusers/models/transformers/auraflow_transformer_2d.py #9715

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f6de4e7
[docs] refactoring docstrings in ./src/diffusers/models/transformers/…
ahnjj Oct 19, 2024
949074f
Update src/diffusers/models/transformers/auraflow_transformer_2d.py
ahnjj Oct 21, 2024
64ff28c
Merge branch 'main' into fix-auraflow_transformer_2d.py
sayakpaul Nov 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 21 additions & 18 deletions src/diffusers/models/transformers/auraflow_transformer_2d.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -187,7 +187,7 @@ class AuraFlowJointTransformerBlock(nn.Module):
* No bias in the attention blocks
* Most LayerNorms are in FP32

Parameters:
Args:
dim (`int`): The number of channels in the input and output.
num_attention_heads (`int`): The number of heads to use for multi-head attention.
attention_head_dim (`int`): The number of channels in each head.
Expand DownExpand Up@@ -257,21 +257,21 @@ class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin):
r"""
A 2D Transformer model as introduced in AuraFlow (https://blog.fal.ai/auraflow/).

Parameters:
Args:
sample_size (`int`): The width of the latent images. This is fixed during training since
it is used to learn a number of position embeddings.
patch_size (`int`): Patch size to turn the input data into small patches.
in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
num_mmdit_layers (`int`, *optional*, defaults to 4): The number of layers of MMDiT Transformer blocks to use.
num_single_dit_layers (`int`, *optional*, defaults to 4):
in_channels (`int`, *optional*, defaults to `16`): The number of channels in the input.
num_mmdit_layers (`int`, *optional*, defaults to `4`): The number of layers of MMDiT Transformer blocks to use.
num_single_dit_layers (`int`, *optional*, defaults to `4`):
The number of layers of Transformer blocks to use. These blocks use concatenated image and text
representations.
attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
attention_head_dim (`int`, *optional*, defaults to `64`): The number of channels in each head.
num_attention_heads (`int`, *optional*, defaults to `18`): The number of heads to use for multi-head attention.
joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
caption_projection_dim (`int`): Number of dimensions to use when projecting the `encoder_hidden_states`.
out_channels (`int`, defaults to 16): Number of output channels.
pos_embed_max_size (`int`, defaults to 4096): Maximum positions to embed from the image latents.
out_channels (`int`, defaults to `16`): Number of output channels.
pos_embed_max_size (`int`, defaults to `4096`): Maximum positions to embed from the image latents.
"""

_no_split_modules = ["AuraFlowJointTransformerBlock", "AuraFlowSingleTransformerBlock", "AuraFlowPatchEmbed"]
Expand DownExpand Up@@ -347,8 +347,9 @@ def __init__(
def attn_processors(self) -> Dict[str, AttentionProcessor]:
r"""
Returns:
`dict` of attention processors: A dictionary containing all attention processors used in the model with
indexed by its weight name.
[`dict[`attention processors`]`]

@a-r-r-o-wa-r-r-o-wOct 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would cause tests to fail too as it requires all implementation that copy from UNet2DConditionalModel.attn_processors to have the same docs and implementation. This can be changed, but you need to make the modification in the mentioned file and then run make fix-copies (can do in a separate PR)

A dictionary containing all attention processors used in the model with
indexed by its weight name.
"""
# set recursively
processors = {}
Expand All@@ -372,13 +373,13 @@ def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, Atte
r"""
Sets the attention processor to use to compute attention.

Parameters:
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
The instantiated processor class or a dictionary of processor classes that will be set as the processor
Args:
processor (`AttentionProcessor` or `Dict[str, AttentionProcessor]`]):
The instantiated processor class or a dictionary of processor classes to be set as the processor
for **all** `Attention` layers.

If `processor` is a dict, the key needs to define the path to the corresponding crossattention
processor. This is strongly recommended when setting trainable attention processors.
If `processor` is a dict, the key needs to define the path to the corresponding cross-attention processor.
This is strongly recommended when setting trainable attention processors.

"""
count = len(self.attn_processors.keys())
Expand All@@ -405,8 +406,10 @@ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedAuraFlowAttnProcessor2_0
def fuse_qkv_projections(self):
"""
Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
are fused. For cross-attention modules, key and value projection matrices are fused.
Enables fused QKV projections.

@a-r-r-o-wa-r-r-o-wOct 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this will cause test to break because it is # Copied from elsewhere. This would require the implementations and docs to be the same as every other occurrence


For self-attention modules, all projection matrices (i.e., query, key, value) are fused.
For cross-attention modules, only key and value projection matrices are fused.

<Tip warning={true}>

Expand Down