refactor rope; diff: 0.14990234375; reason and fix: create rope grid on cpu and move to device

Note: The following line diverges from original behaviour. We create the grid on the device, whereas
original implementation creates it on CPU and then moves it to device. This results in numerical
differences in layerwise debugging outputs, but visually it is the same.

Author: a-r-r-o-w

src/diffusers/models/transformers/transformer_hunyuan_video.py

@@ -24,6 +24,7 @@
 from ...utils import is_torch_version
 from ..attention import FeedForward
 from ..attention_processor import Attention, AttentionProcessor
+from ..embeddings import get_1d_rotary_pos_embed
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
@@ -138,26 +139,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class PatchEmbed(nn.Module):
     def __init__(
         self,
-        patch_size=16,
-        in_chans=3,
-        embed_dim=768,
-        norm_layer=None,
-        flatten=True,
-        bias=True,
-    ):
+        patch_size: Union[int, Tuple[int, int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ) -> None:
         super().__init__()
 
-        patch_size = tuple(patch_size)
-        self.flatten = flatten
-        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+        patch_size = (patch_size, patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
 
-    def forward(self, x):
-        x = self.proj(x)
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        x = self.norm(x)
-        return x
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)  # BCFHW -> BNC
+        return hidden_states
 
 
 class TextProjection(nn.Module):
@@ -384,6 +378,39 @@ def forward(
         return hidden_states
 
 
+class HunyuanVideoRotaryPosEmbed(nn.Module):
+    def __init__(self, patch_size: int, patch_size_t: int, rope_dim: List[int], theta: float = 256.0) -> None:
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.patch_size_t = patch_size_t
+        self.rope_dim = rope_dim
+        self.theta = theta
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size]
+
+        axes_grids = []
+        for i in range(3):
+            # Note: The following line diverges from original behaviour. We create the grid on the device, whereas
+            # original implementation creates it on CPU and then moves it to device. This results in numerical
+            # differences in layerwise debugging outputs, but visually it is the same.
+            grid = torch.arange(0, rope_sizes[i], device=hidden_states.device, dtype=torch.float32)
+            axes_grids.append(grid)
+        grid = torch.meshgrid(*axes_grids, indexing="ij")  # [W, H, T]
+        grid = torch.stack(grid, dim=0)  # [3, W, H, T]
+
+        freqs = []
+        for i in range(3):
+            freq = get_1d_rotary_pos_embed(self.rope_dim[i], grid[i].reshape(-1), self.theta, use_real=True)
+            freqs.append(freq)
+
+        freqs_cos = torch.cat([f[0] for f in freqs], dim=1)  # (W * H * T, D / 2)
+        freqs_sin = torch.cat([f[1] for f in freqs], dim=1)  # (W * H * T, D / 2)
+        return freqs_cos, freqs_sin
+
+
 class HunyuanVideoSingleTransformerBlock(nn.Module):
     def __init__(
         self,
@@ -546,12 +573,12 @@ def __init__(
         guidance_embeds: bool = True,
         text_embed_dim: int = 4096,
         text_embed_dim_2: int = 768,
+        rope_theta: float = 256.0,
     ) -> None:
         super().__init__()
 
         inner_dim = num_attention_heads * attention_head_dim
         out_channels = out_channels or in_channels
-        self.rope_dim_list = rope_dim_list
 
         # image projection
         self.img_in = PatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
@@ -570,6 +597,9 @@ def __init__(
         # guidance modulation
         self.guidance_in = TimestepEmbedder(inner_dim, nn.SiLU)
 
+        # 3. RoPE
+        self.rope = HunyuanVideoRotaryPosEmbed(patch_size, patch_size_t, rope_dim_list, rope_theta)
+
         self.transformer_blocks = nn.ModuleList(
             [
                 HunyuanVideoTransformerBlock(
@@ -664,8 +694,6 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         encoder_attention_mask: torch.Tensor,
         encoder_hidden_states_2: torch.Tensor,
-        freqs_cos: Optional[torch.Tensor] = None,
-        freqs_sin: Optional[torch.Tensor] = None,
         guidance: torch.Tensor = None,
         return_dict: bool = True,
     ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
@@ -676,6 +704,8 @@ def forward(
         post_patch_height = height // p
         post_patch_width = width // p
 
+        image_rotary_emb = self.rope(hidden_states)
+
         temb = self.time_in(timestep)
         temb = temb + self.vector_in(encoder_hidden_states_2)
         temb = temb + self.guidance_in(guidance)
@@ -691,15 +721,14 @@ def forward(
             else lambda x: x
         )
 
-        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
         for _, block in enumerate(self.transformer_blocks):
             hidden_states, encoder_hidden_states = block_forward(block)(
-                hidden_states, encoder_hidden_states, temb, freqs_cis
+                hidden_states, encoder_hidden_states, temb, image_rotary_emb
             )
 
         for block in self.single_transformer_blocks:
             hidden_states, encoder_hidden_states = block_forward(block)(
-                hidden_states, encoder_hidden_states, temb, freqs_cis
+                hidden_states, encoder_hidden_states, temb, image_rotary_emb
             )
 
         hidden_states = self.norm_out(hidden_states, temb)