Merge branch 'main' into b_436918994

Author: entrpn

README.md

@@ -177,11 +177,11 @@ To generate images, run the following command:
   ## LTX-Video
   - In the folder src/maxdiffusion/models/ltx_video/utils, run:
     ```bash
-    python convert_torch_weights_to_jax.py --ckpt_path [LOCAL DIRECTORY FOR WEIGHTS] --transformer_config_path ../xora_v1.2-13B-balanced-128.json
+    python convert_torch_weights_to_jax.py --ckpt_path [LOCAL DIRECTORY FOR WEIGHTS] --transformer_config_path ../ltxv-13B.json
     ```
   - In the repo folder, run:
     ```bash
-    python src/maxdiffusion/generate_ltx_video.py src/maxdiffusion/configs/ltx_video.yml output_dir="[SAME DIRECTORY]" config_path="src/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json"
+    python src/maxdiffusion/generate_ltx_video.py src/maxdiffusion/configs/ltx_video.yml output_dir="[SAME DIRECTORY]" config_path="src/maxdiffusion/models/ltx_video/ltxv-13B.json"
     ```
   - Other generation parameters can be set in ltx_video.yml file.
   ## Flux

…tx_video/xora_v1.2-13B-balanced-128.json …diffusion/models/ltx_video/ltxv-13B.jsonsrc/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json renamed to src/maxdiffusion/models/ltx_video/ltxv-13B.json src/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json renamed to src/maxdiffusion/models/ltx_video/ltxv-13B.json

@@ -60,12 +60,10 @@
 
 def validate_transformer_inputs(prompt_embeds, fractional_coords, latents, encoder_attention_segment_ids):
   # Note: reference shape annotated for first pass default inference parameters
-  max_logging.log("prompts_embeds.shape: ", prompt_embeds.shape, prompt_embeds.dtype)  # (3, 256, 4096) float32
-  max_logging.log("fractional_coords.shape: ", fractional_coords.shape, fractional_coords.dtype)  # (3, 3, 3072) float32
-  max_logging.log("latents.shape: ", latents.shape, latents.dtype)  # (1, 3072, 128) float 32
-  max_logging.log(
-      "encoder_attention_segment_ids.shape: ", encoder_attention_segment_ids.shape, encoder_attention_segment_ids.dtype
-  )  # (3, 256) int32
+  max_logging.log(f"prompts_embeds.shape: {prompt_embeds.shape}")  # (3, 256, 4096) float32
+  max_logging.log(f"fractional_coords.shape: {fractional_coords.shape}")  # (3, 3, 3072) float32
+  max_logging.log(f"latents.shape: {latents.shape}")  # (1, 3072, 128) float 32
+  max_logging.log(f"encoder_attention_segment_ids.shape: {encoder_attention_segment_ids.shape}")  # (3, 256) int32
 
 
 class LTXVideoPipeline:

src/maxdiffusion/pipelines/ltx_video/ltx_video_pipeline.py

@@ -278,7 +278,8 @@ def test_wan_model(self):
       )
     assert dummy_output.shape == hidden_states_shape
 
-  def test_get_qt_provider(self):
+  @patch('maxdiffusion.pipelines.wan.wan_pipeline.qwix.QtRule')
+  def test_get_qt_provider(self, mock_qt_rule):
     """
     Tests the provider logic for all config branches.
     """
@@ -293,9 +294,14 @@ def test_get_qt_provider(self):
     config_int8.quantization = "int8"
     provider_int8: QtProvider = WanPipeline.get_qt_provider(config_int8)
     self.assertIsNotNone(provider_int8)
-    self.assertEqual(provider_int8._rules[0].weight_qtype, jnp.int8)
+    mock_qt_rule.assert_called_once_with(
+      module_path='.*',
+      weight_qtype=jnp.int8,
+      act_qtype=jnp.int8
+    )
 
     # Case 3: Quantization enabled, type 'fp8'
+    mock_qt_rule.reset_mock()
     config_fp8 = Mock(spec=HyperParameters)
     config_fp8.use_qwix_quantization = True
     config_fp8.quantization = "fp8"
@@ -304,6 +310,7 @@ def test_get_qt_provider(self):
     self.assertEqual(provider_fp8.rules[0].kwargs["weight_qtype"], jnp.float8_e4m3fn)
 
     # Case 4: Quantization enabled, type 'fp8_full'
+    mock_qt_rule.reset_mock()
     config_fp8_full = Mock(spec=HyperParameters)
     config_fp8_full.use_qwix_quantization = True
     config_fp8_full.quantization = "fp8_full"
@@ -334,6 +341,8 @@ def test_quantize_transformer_enabled(self, mock_get_dummy_inputs, mock_quantize
     mock_model = Mock(spec=WanModel)
     mock_pipeline = Mock()
     mock_mesh = Mock()
+    mock_mesh.__enter__ = Mock(return_value=None)
+    mock_mesh.__exit__ = Mock(return_value=None)
 
     # Mock the return values of dependencies
     mock_get_dummy_inputs.return_value = (Mock(), Mock(), Mock())