update transformer

Author: a-r-r-o-w

src/diffusers/models/attention_processor.py

@@ -717,144 +717,6 @@ def fuse_projections(self, fuse=True):
         self.fused_projections = fuse
 
 
-class AsymmetricAttention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        query_context_dim: int,
-        num_attentions_heads: int = 8,
-        attention_head_dim: int = 64,
-        bias: bool = False,
-        context_bias: bool = False,
-        out_dim: Optional[int] = None,
-        out_context_dim: Optional[int] = None,
-        qk_norm: Optional[str] = None,
-        eps: float = 1e-5,
-        elementwise_affine: bool = True,
-        out_bias: bool = True,
-        processor: Optional["AttnProcessor"] = None,
-    ) -> None:
-        from .normalization import RMSNorm
-
-        self.query_dim = query_dim
-        self.query_context_dim = query_context_dim
-        self.inner_dim = out_dim if out_dim is not None else num_attentions_heads * attention_head_dim
-        self.out_dim = out_dim if out_dim is not None else query_dim
-        self.out_context_dim = out_context_dim if out_context_dim is not None else query_context_dim
-        
-        self.scale = attention_head_dim ** -0.5
-        self.num_attention_heads = out_dim // attention_head_dim if out_dim is not None else num_attentions_heads
-
-        if qk_norm is None:
-            self.norm_q = None
-            self.norm_k = None
-            self.norm_context_q = None
-            self.norm_context_k = None
-        elif qk_norm == "rms_norm":
-            self.norm_q = RMSNorm(attention_head_dim, eps=eps, elementwise_affine=elementwise_affine)
-            self.norm_k = RMSNorm(attention_head_dim, eps=eps, elementwise_affine=elementwise_affine)
-            self.norm_context_q = RMSNorm(attention_head_dim, eps=eps, elementwise_affine=elementwise_affine)
-            self.norm_context_k = RMSNorm(attention_head_dim, eps=eps, elementwise_affine=elementwise_affine)
-        else:
-            raise ValueError((f"Unknown qk_norm: {qk_norm}. Should be None or `rms_norm`."))
-    
-        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
-        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
-
-        self.to_context_q = nn.Linear(query_context_dim, self.inner_dim, bias=context_bias)
-        self.to_context_k = nn.Linear(query_context_dim, self.inner_dim, bias=context_bias)
-        self.to_context_v = nn.Linear(query_context_dim, self.inner_dim, bias=context_bias)
-
-        # TODO(aryan): Take care of dropouts for training purpose in future
-        self.to_out = nn.ModuleList([
-            nn.Linear(self.inner_dim, self.out_dim)
-        ])
-        self.to_out = nn.ModuleList([
-            nn.Linear(self.inner_dim, self.out_context_dim)
-        ])
-
-        if processor is None:
-            processor = AsymmetricAttnProcessor2_0()
-        
-        self.set_processor(processor)
-
-
-# Similar to SD3
-# class AsymmetricAttnProcessor2_0:
-#     r"""
-#     Processor for implementing Asymmetric SDPA as described in Genmo/Mochi (TODO(aryan) add link).
-#     """
-    
-#     def __init__(self):
-#         if not hasattr(F, "scaled_dot_product_attention"):
-#             raise ImportError("AsymmetricAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-    
-#     def __call__(
-#         self,
-#         attn: AsymmetricAttention,
-#         hidden_states: torch.Tensor,
-#         encoder_hidden_states: torch.Tensor,
-#         temb: torch.Tensor,
-#         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-#     ) -> torch.Tensor:
-#         batch_size = hidden_states.size(0)
-
-#         query = attn.to_q(hidden_states)
-#         key = attn.to_k(hidden_states)
-#         value = attn.to_v(hidden_states)
-
-#         query_context = attn.to_context_q(encoder_hidden_states)
-#         key_context = attn.to_context_k(encoder_hidden_states)
-#         value_context = attn.to_context_v(encoder_hidden_states)
-
-#         inner_dim = key.shape[-1]
-#         head_dim = inner_dim / attn.num_attention_heads
-
-#         query = query.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-#         key = key.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-#         value = value.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-        
-#         query_context = query_context.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-#         key_context = key_context.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-#         value_context = value_context.unflatten(2, (attn.num_attention_heads, head_dim)).transpose(1, 2)
-
-#         if attn.norm_q is not None:
-#             query = attn.norm_q(query)
-#         if attn.norm_k is not None:
-#             key = attn.norm_k(key)
-
-#         if attn.norm_context_q is not None:
-#             query_context = attn.norm_context_q(query_context)
-#             key_context = attn.norm_context_k(key_context)
-        
-#         if image_rotary_emb is not None:
-#             from .embeddings import apply_rotary_emb
-
-#             query = apply_rotary_emb(query, image_rotary_emb)
-#             key = apply_rotary_emb(key, image_rotary_emb)
-
-#         sequence_length = query.size(1)
-#         context_sequence_length = query_context.size(1)
-#         query = torch.cat([query, query_context], dim=1)
-#         key = torch.cat([key, key_context], dim=1)
-#         value = torch.cat([value, value_context], dim=1)
-
-#         hidden_states = F.scaled_dot_product_attention(
-#             query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
-#         )
-
-#         hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
-#         hidden_states = hidden_states.to(query.dtype)
-
-#         hidden_states, encoder_hidden_states = hidden_states.split_with_sizes([sequence_length, context_sequence_length], dim=1)
-        
-#         hidden_states = attn.to_out[0](hidden_states)
-#         encoder_hidden_states = attn.to_context_out[0](encoder_hidden_states)
-
-#         return hidden_states, encoder_hidden_states
-
-
 class AttnProcessor:
     r"""
     Default processor for performing attention-related computations.

src/diffusers/models/embeddings.py

@@ -1302,6 +1302,28 @@ def forward(self, timestep, caption_feat, caption_mask):
         return conditioning
 
 
+class MochiCombinedTimestepCaptionEmbedding(nn.Module):
+    def __init__(self, embedding_dim: int, pooled_projection_dim: int, time_embed_dim: int = 256, num_attention_heads: int = 8) -> None:
+        super().__init__()
+        
+        self.time_proj = Timesteps(
+            num_channels=time_embed_dim, flip_sin_to_cos=True, downscale_freq_shift=0.0
+        )
+        self.timestep_embedder = TimestepEmbedding(in_channels=time_embed_dim, time_embed_dim=embedding_dim)
+        self.pooler = MochiAttentionPool(num_attention_heads=num_attention_heads, embed_dim=pooled_projection_dim, output_dim=embedding_dim)
+        self.caption_proj = nn.Linear(embedding_dim, pooled_projection_dim)
+
+    def forward(self, timestep: torch.LongTensor, encoder_hidden_states: torch.Tensor, encoder_attention_mask: torch.Tensor, hidden_dtype: Optional[torch.dtype] = None):
+        time_proj = self.time_proj(timestep)
+        time_emb = self.timestep_embedder(time_proj.to(dtype=hidden_dtype))
+
+        pooled_projections = self.pooler(encoder_hidden_states, encoder_attention_mask)
+        caption_proj = self.caption_proj(encoder_hidden_states)
+
+        conditioning = time_emb + pooled_projections
+        return conditioning, caption_proj
+
+
 class TextTimeEmbedding(nn.Module):
     def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64):
         super().__init__()

src/diffusers/models/transformers/transformer_mochi.py

@@ -21,11 +21,11 @@
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...utils import logging
 from ...utils.torch_utils import maybe_allow_in_graph
-from ..attention import Attention, FeedForward
-from ..embeddings import PatchEmbed, MochiAttentionPool, TimestepEmbedding, Timesteps
+from ..attention import Attention, FeedForward, JointAttnProcessor2_0
+from ..embeddings import PatchEmbed, MochiCombinedTimestepCaptionEmbedding
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
-from ..normalization import MochiRMSNormZero, RMSNorm
+from ..normalization import AdaLayerNormContinuous, MochiRMSNormZero, RMSNorm
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
@@ -38,61 +38,73 @@ def __init__(
         dim: int,
         num_attention_heads: int,
         attention_head_dim: int,
-        caption_dim: int,
+        pooled_projection_dim: int,
+        qk_norm: str = "rms_norm",
         activation_fn: str = "swiglu",
-        update_captions: bool = True,
+        context_pre_only: bool = True,
     ) -> None:
         super().__init__()
 
-        self.update_captions = update_captions
+        self.context_pre_only = context_pre_only
 
         self.norm1 = MochiRMSNormZero(dim, 4 * dim)
 
-        if update_captions:
-            self.norm_context1 = MochiRMSNormZero(dim, 4 * caption_dim)
+        if context_pre_only:
+            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim)
         else:
-            self.norm_context1 = RMSNorm(caption_dim, eps=1e-5, elementwise_affine=False)
+            self.norm1_context = RMSNorm(pooled_projection_dim, eps=1e-6, elementwise_affine=False)
 
         self.attn = Attention(
             query_dim=dim,
             heads=num_attention_heads,
             attention_head_dim=attention_head_dim,
             out_dim=4 * dim,
-            qk_norm="rms_norm",
-            eps=1e-5,
-            elementwise_affine=False,
-        )
-        self.attn_context = Attention(
-            query_dim=dim,
-            heads=num_attention_heads,
-            attention_head_dim=attention_head_dim,
-            out_dim=4 * caption_dim if update_captions else caption_dim,
-            qk_norm="rms_norm",
-            eps=1e-5,
+            qk_norm=qk_norm,
+            eps=1e-6,
             elementwise_affine=False,
+            processor=JointAttnProcessor2_0(),
         )
 
+        self.norm2 = RMSNorm(dim, eps=1e-6, elementwise_affine=False)
+        self.norm2_context = RMSNorm(pooled_projection_dim, eps=1e-6, elementwise_affine=False)
+        
+        self.norm3 = RMSNorm(dim, eps=1e-6, elementwise_affine=False)
+        self.norm3_context = RMSNorm(pooled_projection_dim, eps=1e-56, elementwise_affine=False)
+
         self.ff = FeedForward(dim, mult=4, activation_fn=activation_fn)
-        self.ff_context = FeedForward(caption_dim, mult=4, activation_fn=activation_fn)
+        self.ff_context = FeedForward(pooled_projection_dim, mult=4, activation_fn=activation_fn)
+
+        self.norm4 = RMSNorm(dim, eps=1e-6, elementwise_affine=False)
+        self.norm4_context = RMSNorm(pooled_projection_dim, eps=1e-56, elementwise_affine=False)
 
     def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, image_rotary_emb: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
         norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
 
-        if self.update_captions:
-            norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm_context1(encoder_hidden_states, temb)
+        if self.context_pre_only:
+            norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context(encoder_hidden_states, temb)
         else:
-            norm_encoder_hidden_states = self.norm_context1(encoder_hidden_states)
+            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states)
 
-        attn_hidden_states = self.attn(
+        attn_hidden_states, context_attn_hidden_states = self.attn(
             hidden_states=norm_hidden_states,
-            encoder_hidden_states=None,
+            encoder_hidden_states=norm_encoder_hidden_states,
             image_rotary_emb=image_rotary_emb,
         )
-        attn_encoder_hidden_states = self.attn_context(
-            hidden_states=norm_encoder_hidden_states,
-            encoder_hidden_states=None,
-            image_rotary_emb=None,
-        )
+
+        hidden_states = hidden_states + self.norm2(attn_hidden_states) * torch.tanh(gate_msa).unsqueeze(1)
+        hidden_states = self.norm3(hidden_states) * (1 + scale_mlp.unsqueeze(1))
+        if not self.context_pre_only:
+            encoder_hidden_states = encoder_hidden_states + self.norm2_context(context_attn_hidden_states) * torch.tanh(enc_gate_msa).unsqueeze(1)
+            encoder_hidden_states = encoder_hidden_states + self.norm3_context(encoder_hidden_states) * (1 + enc_scale_mlp.unsqueeze(1))
+        
+        ff_output = self.ff(hidden_states)
+        context_ff_output = self.ff_context(encoder_hidden_states)
+        
+        hidden_states = hidden_states + ff_output * torch.tanh(gate_mlp).unsqueeze(1)
+        if not self.context_pre_only:
+            encoder_hidden_states = encoder_hidden_states + context_ff_output * torch.tanh(enc_gate_mlp).unsqueeze(0)
+        
+        return hidden_states, encoder_hidden_states
 
 
 @maybe_allow_in_graph
@@ -106,32 +118,35 @@ def __init__(
         num_attention_heads: int = 24,
         attention_head_dim: int = 128,
         num_layers: int = 48,
-        caption_dim=1536,
-        mlp_ratio_x=4.0,
-        mlp_ratio_y=4.0,
+        pooled_projection_dim: int = 1536,
         in_channels=12,
-        qk_norm=True,
-        qkv_bias=False,
-        out_bias=True,
+        out_channels: Optional[int] = None,
+        qk_norm: str = "rms_norm",
         timestep_mlp_bias=True,
         timestep_scale=1000.0,
-        text_embed_dim=4096,
+        text_embed_dim: int = 4096,
+        time_embed_dim: int = 256,
         activation_fn: str = "swiglu",
-        max_sequence_length=256,
+        max_sequence_length: int = 256,
     ) -> None:
         super().__init__()
 
         inner_dim = num_attention_heads * attention_head_dim
+        out_channels = out_channels or in_channels
+
+        self.time_embed = MochiCombinedTimestepCaptionEmbedding(
+            embedding_dim=text_embed_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            time_embed_dim=time_embed_dim,
+            num_attention_heads=8,
+        )
 
         self.patch_embed = PatchEmbed(
             patch_size=patch_size,
             in_channels=in_channels,
             embed_dim=inner_dim,
         )
 
-        self.caption_embedder = MochiAttentionPool(num_attention_heads=8, embed_dim=text_embed_dim, output_dim=inner_dim)
-        self.caption_proj = nn.Linear(text_embed_dim, caption_dim)
-
         self.pos_frequencies = nn.Parameter(
             torch.empty(3, num_attention_heads, attention_head_dim // 2)
         )
@@ -141,9 +156,53 @@ def __init__(
                 dim=inner_dim,
                 num_attention_heads=num_attention_heads,
                 attention_head_dim=attention_head_dim,
-                caption_dim=caption_dim,
+                pooled_projection_dim=pooled_projection_dim,
+                qk_norm=qk_norm,
                 activation_fn=activation_fn,
-                update_captions=i < num_layers - 1,
+                context_pre_only=i < num_layers - 1,
             )
             for i in range(num_layers)
         ])
+
+        self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6, norm_type="layer_norm")
+        self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
+    
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        encoder_attention_mask: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        return_dict: bool = True,
+    ) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p = self.config.patch_size
+
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        temb, caption_proj = self.time_embed(timestep, encoder_hidden_states, encoder_attention_mask)
+
+        hidden_states = self.patch_embed(hidden_states)
+
+        for i, block in enumerate(self.transformer_blocks):
+            hidden_states, encoder_hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+            )
+        
+        # TODO(aryan): do something with self.pos_frequencies
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_height, p, p, -1)
+        hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
+        output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)