reformatted

Author: prishajain1

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -25,6 +25,7 @@
 DType = common_types.DType
 BlockSizes = common_types.BlockSizes
 
+
 def apply_rotary_emb(x: Array, freqs: Tuple[Array, Array]) -> Array:
   """
   Applies Interleaved RoPE to input x.
@@ -193,9 +194,7 @@ def prepare_video_coords(
     # pixel_coords[:, 0, ...] selects Frame dimension.
     # pixel_coords shape: [B, 3, num_patches, 2] -> dim 1 is (F, H, W)
     frame_coords = pixel_coords[:, 0, ...]
-    frame_coords = jnp.clip(
-        frame_coords + self.causal_offset - self.scale_factors[0], min=0
-    )
+    frame_coords = jnp.clip(frame_coords + self.causal_offset - self.scale_factors[0], min=0)
     pixel_coords = pixel_coords.at[:, 0, ...].set(frame_coords / fps)
 
     return pixel_coords
@@ -212,16 +211,12 @@ def prepare_audio_coords(
     # 2. Start timestamps
     audio_scale_factor = self.scale_factors[0]
     grid_start_mel = grid_f * audio_scale_factor
-    grid_start_mel = jnp.clip(
-        grid_start_mel + self.causal_offset - audio_scale_factor, min=0
-    )
+    grid_start_mel = jnp.clip(grid_start_mel + self.causal_offset - audio_scale_factor, min=0)
     grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate
 
     # 3. End timestamps
     grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor
-    grid_end_mel = jnp.clip(
-        grid_end_mel + self.causal_offset - audio_scale_factor, min=0
-    )
+    grid_end_mel = jnp.clip(grid_end_mel + self.causal_offset - audio_scale_factor, min=0)
     grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate
 
     # Stack [num_patches, 2]