fixes

Author: a-r-r-o-w

src/diffusers/models/attention_processor.py

@@ -1795,8 +1795,7 @@ def __call__(
             # dropout
             hidden_states = attn.to_out[1](hidden_states)
 
-            if hasattr(attn, "to_add_out"):
-                encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
 
             return hidden_states, encoder_hidden_states
         else:
@@ -3082,6 +3081,89 @@ def __call__(
         return hidden_states
 
 
+class MochiAttnProcessor2_0:
+    """Attention processor used in Mochi."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        breakpoint()
+        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 = query.unflatten(2, (attn.heads, -1))
+        key = key.unflatten(2, (attn.heads, -1))
+        value = value.unflatten(2, (attn.heads, -1))
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        encoder_query = attn.add_q_proj(encoder_hidden_states)
+        encoder_key = attn.add_k_proj(encoder_hidden_states)
+        encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+        encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
+        encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
+        encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
+
+        if attn.norm_added_q is not None:
+            encoder_query = attn.norm_added_q(encoder_query)
+        if attn.norm_added_k is not None:
+            encoder_key = attn.norm_added_k(encoder_key)
+        
+        if image_rotary_emb is not None:
+            def apply_rotary_emb(x, freqs_cos, freqs_sin):
+                x_even = x[..., 0::2].float()
+                x_odd = x[..., 1::2].float()
+
+                cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
+                sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
+
+                return torch.stack([cos, sin], dim=-1).flatten(-2)
+
+            query = apply_rotary_emb(query, *image_rotary_emb)
+            key = apply_rotary_emb(key, *image_rotary_emb)
+        
+        query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+        encoder_query, encoder_key, encoder_value = encoder_query.transpose(1, 2), encoder_key.transpose(1, 2), encoder_value.transpose(1, 2)
+
+        sequence_length = query.size(2)
+        encoder_sequence_length = encoder_query.size(2)
+
+        query = torch.cat([query, encoder_query], dim=2)
+        key = torch.cat([key, encoder_key], dim=2)
+        value = torch.cat([value, encoder_value], dim=2)
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, 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, encoder_sequence_length), dim=1)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+
+
 class FusedAttnProcessor2_0:
     r"""
     Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). It uses

src/diffusers/models/normalization.py

@@ -246,13 +246,13 @@ class MochiRMSNormZero(nn.Module):
     """
 
     def __init__(
-        self, embedding_dim: int, hidden_dim: int, norm_eps: float = 1e-5, elementwise_affine: bool = False
+        self, embedding_dim: int, hidden_dim: int, 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)
+        self.norm = RMSNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
 
     def forward(
         self, hidden_states: torch.Tensor, emb: torch.Tensor

src/diffusers/models/transformers/transformer_mochi.py

@@ -22,7 +22,7 @@
 from ...utils import logging
 from ...utils.torch_utils import maybe_allow_in_graph
 from ..attention import FeedForward
-from ..attention_processor import Attention, FluxAttnProcessor2_0
+from ..attention_processor import Attention, MochiAttnProcessor2_0
 from ..embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
@@ -43,22 +43,23 @@ def __init__(
         qk_norm: str = "rms_norm",
         activation_fn: str = "swiglu",
         context_pre_only: bool = True,
+        eps: float = 1e-6,
     ) -> None:
         super().__init__()
 
         self.context_pre_only = context_pre_only
         self.ff_inner_dim = (4 * dim * 2) // 3
         self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3
 
-        self.norm1 = MochiRMSNormZero(dim, 4 * dim)
+        self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
 
         if not context_pre_only:
-            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim)
+            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
         else:
             self.norm1_context = LuminaLayerNormContinuous(
                 embedding_dim=pooled_projection_dim,
                 conditioning_embedding_dim=dim,
-                eps=1e-6,
+                eps=eps,
                 elementwise_affine=False,
                 norm_type="rms_norm",
                 out_dim=None,
@@ -76,16 +77,16 @@ def __init__(
             out_dim=dim,
             out_context_dim=pooled_projection_dim,
             context_pre_only=context_pre_only,
-            processor=FluxAttnProcessor2_0(),
-            eps=1e-6,
+            processor=MochiAttnProcessor2_0(),
+            eps=eps,
             elementwise_affine=True,
         )
 
-        self.norm2 = RMSNorm(dim, eps=1e-6, elementwise_affine=False)
-        self.norm2_context = RMSNorm(pooled_projection_dim, eps=1e-6, elementwise_affine=False)
+        self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2_context = RMSNorm(pooled_projection_dim, eps=eps, 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.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
 
         self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
         self.ff_context = None
@@ -94,8 +95,8 @@ def __init__(
                 pooled_projection_dim, inner_dim=self.ff_context_inner_dim, activation_fn=activation_fn, bias=False
             )
 
-        self.norm4 = RMSNorm(dim, eps=1e-6, elementwise_affine=False)
-        self.norm4_context = RMSNorm(pooled_projection_dim, eps=1e-56, elementwise_affine=False)
+        self.norm4 = RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm4_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
 
     def forward(
         self,
@@ -104,6 +105,7 @@ def forward(
         temb: torch.Tensor,
         image_rotary_emb: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        breakpoint()
         norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
 
         if not self.context_pre_only:
@@ -140,6 +142,40 @@ def forward(
         return hidden_states, encoder_hidden_states
 
 
+class MochiRoPE(nn.Module):
+    def __init__(self, base_height: int = 192, base_width: int = 192, theta: float = 10000.0) -> None:
+        super().__init__()
+        
+        self.target_area = base_height * base_width
+    
+    def _centers(self, start, stop, num, device, dtype) -> torch.Tensor:
+        edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype)
+        return (edges[:-1] + edges[1:]) / 2
+    
+    def _get_positions(self, num_frames: int, height: int, width: int, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+        scale = (self.target_area / (height * width)) ** 0.5
+        
+        t = torch.arange(num_frames, device=device, dtype=dtype)
+        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
+        w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
+
+        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
+
+        positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
+        return positions
+
+    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        freqs = torch.einsum("nd,dhf->nhf", pos, freqs)
+        freqs_cos = torch.cos(freqs)
+        freqs_sin = torch.sin(freqs)
+        return freqs_cos, freqs_sin
+
+    def forward(self, pos_frequencies: torch.Tensor, num_frames: int, height: int, width: int, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        pos = self._get_positions(num_frames, height, width, device, dtype)
+        rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
+        return rope_cos, rope_sin
+
+
 @maybe_allow_in_graph
 class MochiTransformer3DModel(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
@@ -169,6 +205,7 @@ def __init__(
             patch_size=patch_size,
             in_channels=in_channels,
             embed_dim=inner_dim,
+            pos_embed_type=None,
         )
 
         self.time_embed = MochiCombinedTimestepCaptionEmbedding(
@@ -180,6 +217,7 @@ def __init__(
         )
 
         self.pos_frequencies = nn.Parameter(torch.empty(3, num_attention_heads, attention_head_dim // 2))
+        self.rope = MochiRoPE()
 
         self.transformer_blocks = nn.ModuleList(
             [
@@ -207,7 +245,6 @@ def forward(
         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
@@ -224,6 +261,8 @@ def forward(
         hidden_states = self.patch_embed(hidden_states)
         hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
 
+        image_rotary_emb = self.rope(self.pos_frequencies, num_frames, post_patch_height, post_patch_width, device=hidden_states.device, dtype=torch.float32)
+
         for i, block in enumerate(self.transformer_blocks):
             hidden_states, encoder_hidden_states = block(
                 hidden_states=hidden_states,