feat: add LTX2 smoke test and update setup.sh environment

Author: Perseus14

src/maxdiffusion/tests/ltx2_smoke_test.py

@@ -17,7 +17,6 @@
 import os
 import time
 import unittest
-import pytest
 import jax
 import jax.numpy as jnp
 
@@ -32,6 +31,7 @@
 IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true"
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 
+
 class LTX2SmokeTest(unittest.TestCase):
   """End-to-end smoke test for LTX2."""
 
@@ -43,9 +43,9 @@ def setUpClass(cls):
             None,
             os.path.join(THIS_DIR, "..", "configs", "ltx2_video.yml"),
             "num_inference_steps=2",  # Small number of steps for fast test
-            "height=256",             # Small resolution
+            "height=256",  # Small resolution
             "width=256",
-            "num_frames=9",           # Small number of frames
+            "num_frames=9",  # Small number of frames
             "seed=0",
             "attention=flash",
         ],
@@ -62,7 +62,7 @@ def setUpClass(cls):
   def test_ltx2_inference(self):
     """Test that LTX2 pipeline can run inference and produce output."""
     generator = jax.random.key(self.config.seed)
-    
+
     t0 = time.perf_counter()
     out = self.pipeline(
         prompt=self.prompt,
@@ -76,22 +76,22 @@ def test_ltx2_inference(self):
         dtype=jnp.bfloat16,
     )
     t1 = time.perf_counter()
-    
+
     print(f"LTX2 Inference took: {t1 - t0:.2f}s")
-    
+
     videos = out.frames if hasattr(out, "frames") else out[0]
     audios = out.audio if hasattr(out, "audio") else None
-    
+
     self.assertIsNotNone(videos)
     # Check that we got frames
     self.assertGreater(len(videos), 0)
-    
+
     # LTX2 might also produce audio, check if it's there if expected
     # The config doesn't explicitly say if it's T2AV or just T2V, but the pipeline seems to handle audio.
     # We can just log if audio is present.
     if audios is not None:
-        print(f"Audio produced with shape: {audios[0].shape}")
-        self.assertGreater(len(audios), 0)
+      print(f"Audio produced with shape: {audios[0].shape}")
+      self.assertGreater(len(audios), 0)
 
 if __name__ == "__main__":
   unittest.main()