lint

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1346,9 +1346,9 @@ def __call__(
         audio_embeds_sharded = jax.device_put(audio_embeds, spec)
 
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
-      
+
       scan_diffusion_loop = getattr(self.config, "scan_diffusion_loop", True)
-      
+
       if scan_diffusion_loop:
         latents_jax, audio_latents_jax = run_diffusion_loop(
             graphdef,
@@ -1375,7 +1375,7 @@ def __call__(
         # Old Python loop path
         latents_jax = latents_jax.astype(jnp.float32)
         audio_latents_jax = audio_latents_jax.astype(jnp.float32)
-        
+
         for t in timesteps_jax:
           noise_pred, noise_pred_audio = transformer_forward_pass(
               graphdef,
@@ -1395,26 +1395,28 @@ def __call__(
               audio_num_frames,
               frame_rate,
           )
-          
+
           if guidance_scale > 1.0:
             noise_pred_uncond, noise_pred_text = jnp.split(noise_pred, 2, axis=0)
             noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-            
+
             noise_pred_audio_uncond, noise_pred_audio_text = jnp.split(noise_pred_audio, 2, axis=0)
             noise_pred_audio = noise_pred_audio_uncond + guidance_scale * (noise_pred_audio_text - noise_pred_audio_uncond)
-            
+
             latents_step = latents_jax[batch_size:]
             audio_latents_step = audio_latents_jax[batch_size:]
           else:
             latents_step = latents_jax
             audio_latents_step = audio_latents_jax
-            
+
           latents_step, _ = self.scheduler.step(scheduler_state, noise_pred, t, latents_step, return_dict=False)
           latents_step = latents_step.astype(jnp.float32)
-          
-          audio_latents_step, _ = self.scheduler.step(scheduler_state, noise_pred_audio, t, audio_latents_step, return_dict=False)
+
+          audio_latents_step, _ = self.scheduler.step(
+              scheduler_state, noise_pred_audio, t, audio_latents_step, return_dict=False
+          )
           audio_latents_step = audio_latents_step.astype(jnp.float32)
-          
+
           if guidance_scale > 1.0:
             latents_jax = jnp.concatenate([latents_step] * 2, axis=0)
             audio_latents_jax = jnp.concatenate([audio_latents_step] * 2, axis=0)