reformatted

Author: prishajain1

src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2.py

@@ -225,7 +225,9 @@ def __init__(
       self.per_channel_scale2 = None
 
     if timestep_conditioning:
-      self.scale_shift_table = nnx.Param(jax.random.normal(rngs.params(), (4, in_channels), dtype=dtype) / (in_channels**0.5))
+      self.scale_shift_table = nnx.Param(
+          jax.random.normal(rngs.params(), (4, in_channels), dtype=dtype) / (in_channels**0.5)
+      )
     else:
       self.scale_shift_table = None
 
@@ -1261,40 +1263,37 @@ def enable_tiling(
   def blend_v(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[2], b.shape[2], blend_extent)
     if blend_extent <= 0:
-        return b
-        
+      return b
+
     # Create broadcastable blending weights: (1, 1, blend_extent, 1, 1)
     y = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, 1, -1, 1, 1)
-    
-    val = a[:, :, -blend_extent:, :, :] * (1.0 - y / blend_extent) + \
-          b[:, :, :blend_extent, :, :] * (y / blend_extent)
-          
+
+    val = a[:, :, -blend_extent:, :, :] * (1.0 - y / blend_extent) + b[:, :, :blend_extent, :, :] * (y / blend_extent)
+
     return b.at[:, :, :blend_extent, :, :].set(val)
 
   def blend_h(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[3], b.shape[3], blend_extent)
     if blend_extent <= 0:
-        return b
-        
+      return b
+
     # Create broadcastable blending weights: (1, 1, 1, blend_extent, 1)
     x = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, 1, 1, -1, 1)
-    
-    val = a[:, :, :, -blend_extent:, :] * (1.0 - x / blend_extent) + \
-          b[:, :, :, :blend_extent, :] * (x / blend_extent)
-          
+
+    val = a[:, :, :, -blend_extent:, :] * (1.0 - x / blend_extent) + b[:, :, :, :blend_extent, :] * (x / blend_extent)
+
     return b.at[:, :, :, :blend_extent, :].set(val)
 
   def blend_t(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[1], b.shape[1], blend_extent)
     if blend_extent <= 0:
-        return b
-        
+      return b
+
     # Create broadcastable blending weights: (1, blend_extent, 1, 1, 1)
     x = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, -1, 1, 1, 1)
-    
-    val = a[:, -blend_extent:, :, :, :] * (1.0 - x / blend_extent) + \
-          b[:, :blend_extent, :, :, :] * (x / blend_extent)
-          
+
+    val = a[:, -blend_extent:, :, :, :] * (1.0 - x / blend_extent) + b[:, :blend_extent, :, :, :] * (x / blend_extent)
+
     return b.at[:, :blend_extent, :, :, :].set(val)
 
   def tiled_encode(self, x: jax.Array, key: Optional[jax.Array] = None, causal: Optional[bool] = None) -> jax.Array: