make style

Author: a-r-r-o-w

src/diffusers/models/embeddings.py

@@ -1303,17 +1303,25 @@ def forward(self, timestep, caption_feat, caption_mask):
 
 
 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:
+    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.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.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):
+    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))
 
@@ -1467,7 +1475,7 @@ def __init__(
         self.to_kv = nn.Linear(embed_dim, 2 * embed_dim)
         self.to_q = nn.Linear(embed_dim, embed_dim)
         self.to_out = nn.Linear(embed_dim, self.output_dim)
-    
+
     @staticmethod
     def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor:
         """
@@ -1526,9 +1534,7 @@ def forward(self, x: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:
         q = q.unsqueeze(2)  # (B, H, 1, head_dim)
 
         # Compute attention.
-        x = F.scaled_dot_product_attention(
-            q, k, v, attn_mask=attn_mask, dropout_p=0.0
-        )  # (B, H, 1, head_dim)
+        x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)  # (B, H, 1, head_dim)
 
         # Concatenate heads and run output.
         x = x.squeeze(2).flatten(1, 2)  # (B, D = H * head_dim)

src/diffusers/models/normalization.py

@@ -245,14 +245,18 @@ class MochiRMSNormZero(nn.Module):
         embedding_dim (`int`): The size of each embedding vector.
     """
 
-    def __init__(self, embedding_dim: int, hidden_dim: int, norm_eps: float = 1e-5, elementwise_affine: bool = False) -> None:
+    def __init__(
+        self, embedding_dim: int, hidden_dim: int, norm_eps: float = 1e-5, elementwise_affine: bool = False
+    ) -> None:
         super().__init__()
 
         self.silu = nn.SiLU()
         self.linear = nn.Linear(embedding_dim, hidden_dim)
         self.norm = RMSNorm(embedding_dim, eps=norm_eps, elementwise_affine=elementwise_affine)
 
-    def forward(self, hidden_states: torch.Tensor, emb: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def forward(
+        self, hidden_states: torch.Tensor, emb: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         emb = self.linear(self.silu(emb))
         scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
         hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None])

src/diffusers/models/transformers/transformer_mochi.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -22,7 +22,7 @@
 from ...utils import logging
 from ...utils.torch_utils import maybe_allow_in_graph
 from ..attention import Attention, FeedForward, JointAttnProcessor2_0
-from ..embeddings import PatchEmbed, MochiCombinedTimestepCaptionEmbedding
+from ..embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, MochiRMSNormZero, RMSNorm
@@ -46,14 +46,14 @@ def __init__(
         super().__init__()
 
         self.context_pre_only = context_pre_only
-        
+
         self.norm1 = MochiRMSNormZero(dim, 4 * dim)
 
         if context_pre_only:
             self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim)
         else:
             self.norm1_context = RMSNorm(pooled_projection_dim, eps=1e-6, elementwise_affine=False)
-        
+
         self.attn = Attention(
             query_dim=dim,
             heads=num_attention_heads,
@@ -67,7 +67,7 @@ def __init__(
 
         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)
 
@@ -76,15 +76,23 @@ def __init__(
 
         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]:
+
+    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.context_pre_only:
-            norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context(encoder_hidden_states, temb)
+            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.norm1_context(encoder_hidden_states)
-        
+
         attn_hidden_states, context_attn_hidden_states = self.attn(
             hidden_states=norm_hidden_states,
             encoder_hidden_states=norm_encoder_hidden_states,
@@ -94,16 +102,20 @@ def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tens
         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))
-        
+            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
 
 
@@ -140,33 +152,35 @@ def __init__(
             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.pos_frequencies = nn.Parameter(
-            torch.empty(3, num_attention_heads, attention_head_dim // 2)
+        self.pos_frequencies = nn.Parameter(torch.empty(3, num_attention_heads, attention_head_dim // 2))
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MochiTransformerBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    pooled_projection_dim=pooled_projection_dim,
+                    qk_norm=qk_norm,
+                    activation_fn=activation_fn,
+                    context_pre_only=i < num_layers - 1,
+                )
+                for i in range(num_layers)
+            ]
         )
 
-        self.transformer_blocks = nn.ModuleList([
-            MochiTransformerBlock(
-                dim=inner_dim,
-                num_attention_heads=num_attention_heads,
-                attention_head_dim=attention_head_dim,
-                pooled_projection_dim=pooled_projection_dim,
-                qk_norm=qk_norm,
-                activation_fn=activation_fn,
-                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.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,
@@ -193,13 +207,13 @@ def forward(
                 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.reshape(batch_size, num_frames, post_patch_height, post_patch_width, 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)