fix positional encoding from becoming > 4096

Author: kashif

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -235,7 +235,7 @@ def forward(
         video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
         txt_seq_lens: Optional[List[int]] = None,
         device: torch.device = None,
-        max_txt_seq_len: Optional[Union[int, torch.Tensor]] = None,
+        max_txt_seq_len: Optional[int] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Args:
@@ -245,9 +245,9 @@ def forward(
                 Deprecated parameter. Use `max_txt_seq_len` instead. If provided, the maximum value will be used.
             device: (`torch.device`, *optional*):
                 The device on which to perform the RoPE computation.
-            max_txt_seq_len (`int` or `torch.Tensor`, *optional*):
+            max_txt_seq_len (`int`, *optional*):
                 The maximum text sequence length for RoPE computation. This should match the encoder hidden states
-                sequence length. Can be either an int or a scalar tensor (for torch.compile compatibility).
+                sequence length.
         """
         # Handle deprecated txt_seq_lens parameter
         if txt_seq_lens is not None:
@@ -296,9 +296,14 @@ def forward(
             else:
                 max_vid_index = max(height, width, max_vid_index)
 
-        max_txt_seq_len_int = int(max_txt_seq_len)
-        # Create device-specific copy for text freqs without modifying self.pos_freqs
-        txt_freqs = self.pos_freqs.to(device)[max_vid_index : max_vid_index + max_txt_seq_len_int, ...]
+        pos_freqs = self.pos_freqs.to(device)
+
+        # Clamp text sequence length to avoid buffer overflow
+        buffer_size = pos_freqs.shape[0]
+        available_space = buffer_size - max_vid_index
+        safe_txt_seq_len = min(max_txt_seq_len, available_space)
+
+        txt_freqs = pos_freqs[max_vid_index : max_vid_index + safe_txt_seq_len]
         vid_freqs = torch.cat(vid_freqs, dim=0)
 
         return vid_freqs, txt_freqs
@@ -367,17 +372,17 @@ def rope_params(self, index, dim, theta=10000):
     def forward(
         self,
         video_fhw: Union[Tuple[int, int, int], List[Tuple[int, int, int]]],
-        max_txt_seq_len: Union[int, torch.Tensor],
+        max_txt_seq_len: int,
         device: torch.device = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Args:
             video_fhw (`Tuple[int, int, int]` or `List[Tuple[int, int, int]]`):
                 A list of 3 integers [frame, height, width] representing the shape of the video, or a list of layer
                 structures.
-            max_txt_seq_len (`int` or `torch.Tensor`):
+            max_txt_seq_len (`int`):
                 The maximum text sequence length for RoPE computation. This should match the encoder hidden states
-                sequence length. Can be either an int or a scalar tensor (for torch.compile compatibility).
+                sequence length.
             device: (`torch.device`, *optional*):
                 The device on which to perform the RoPE computation.
         """
@@ -417,9 +422,15 @@ def forward(
                 max_vid_index = max(height, width, max_vid_index)
 
         max_vid_index = max(max_vid_index, layer_num)
-        max_txt_seq_len_int = int(max_txt_seq_len)
-        # Create device-specific copy for text freqs without modifying self.pos_freqs
-        txt_freqs = self.pos_freqs.to(device)[max_vid_index : max_vid_index + max_txt_seq_len_int, ...]
+
+        pos_freqs = self.pos_freqs.to(device)
+
+        # Clamp text sequence length to avoid buffer overflow
+        buffer_size = pos_freqs.shape[0]
+        available_space = buffer_size - max_vid_index
+        safe_txt_seq_len = min(max_txt_seq_len, available_space)
+
+        txt_freqs = pos_freqs[max_vid_index : max_vid_index + safe_txt_seq_len]
         vid_freqs = torch.cat(vid_freqs, dim=0)
 
         return vid_freqs, txt_freqs
@@ -920,7 +931,7 @@ def forward(
         encoder_hidden_states = self.txt_in(encoder_hidden_states)
 
         # Use the encoder_hidden_states sequence length for RoPE computation and normalize mask
-        text_seq_len, _, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
+        text_seq_len, per_sample_len, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
             encoder_hidden_states, encoder_hidden_states_mask
         )
 
@@ -933,6 +944,8 @@ def forward(
             else self.time_text_embed(timestep, guidance, hidden_states, additional_t_cond)
         )
 
+        # Pass the static text_seq_len to RoPE (encoder_hidden_states.shape[1])
+        # The RoPE class will clamp it to avoid buffer overflow
         image_rotary_emb = self.pos_embed(img_shapes, max_txt_seq_len=text_seq_len, device=hidden_states.device)
 
         block_attention_kwargs = attention_kwargs.copy() if attention_kwargs is not None else {}