ltx2 pipeline

Author: prishajain1

src/maxdiffusion/tests/ltx2_pipeline_test.py

@@ -172,5 +172,38 @@ def test_pipeline_call(self):
         # Let's inspect output shape in actual running test
         print(f"Output shape: {output.shape}")
 
+    def test_pipeline_call_i2v(self):
+        pipeline = LTX2Pipeline(
+            tokenizer=self.tokenizer,
+            text_encoder=self.text_encoder,
+            text_encoder_connector=self.text_encoder_connector,
+            transformer=self.transformer,
+            vae=self.vae,
+            scheduler=self.scheduler,
+            scheduler_state=self.scheduler_state,
+            devices_array=self.devices_array,
+            mesh=self.mesh,
+            config=self.config
+        )
+        
+        # Create dummy image
+        batch = 1
+        height = 32
+        width = 32
+        # Image shape: (B, H, W, C)
+        image = np.random.rand(batch, height, width, 3).astype(np.float32)
+        
+        output = pipeline(
+            prompt="A cat dancing",
+            image=image,
+            height=32,
+            width=32,
+            num_frames=9,
+            num_inference_steps=1,
+            num_videos_per_prompt=1
+        )
+        self.assertIsNotNone(output)
+        print(f"I2V Output shape: {output.shape}")
+
 if __name__ == '__main__':
     unittest.main()