Merge remote-tracking branch 'origin/main' into prisha/ltx2_transformer

Author: prishajain1

src/maxdiffusion/generate.py

@@ -26,6 +26,7 @@
 from absl import app
 from maxdiffusion import (pyconfig, FlaxDDIMScheduler, max_utils)
 
+from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.maxdiffusion_utils import rescale_noise_cfg
 from flax.linen import partitioning as nn_partitioning
 from maxdiffusion.image_processor import VaeImageProcessor
@@ -261,4 +262,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/generate_flux.py

@@ -33,6 +33,7 @@
 
 from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.max_utils import (
     device_put_replicated,
     get_memory_allocations,
@@ -492,4 +493,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/generate_flux_pipeline.py

@@ -26,6 +26,7 @@
 from maxdiffusion import pyconfig, max_logging, max_utils
 
 from maxdiffusion.checkpointing.checkpointing_utils import load_params_from_path
+from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.max_utils import setup_initial_state
 
 
@@ -123,4 +124,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/generate_ltx_video.py

@@ -21,6 +21,7 @@
 from maxdiffusion.pipelines.ltx_video.ltx_video_pipeline import LTXMultiScalePipeline, ConditioningItem
 import maxdiffusion.pipelines.ltx_video.crf_compressor as crf_compressor
 from maxdiffusion import pyconfig, max_logging
+from maxdiffusion.train_utils import transformer_engine_context
 import torchvision.transforms.functional as TVF
 import imageio
 from datetime import datetime
@@ -267,4 +268,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/generate_sdxl.py

@@ -29,6 +29,7 @@
 
 from maxdiffusion import pyconfig, max_utils
 from maxdiffusion.image_processor import VaeImageProcessor
+from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.maxdiffusion_utils import (
     get_add_time_ids,
     rescale_noise_cfg,
@@ -322,4 +323,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/generate_wan.py

@@ -23,6 +23,7 @@
 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p2 import WanCheckpointerI2V_2_2
 from maxdiffusion import pyconfig, max_logging, max_utils
 from absl import app
+from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.utils import export_to_video
 from maxdiffusion.utils.loading_utils import load_image
 from google.cloud import storage
@@ -296,4 +297,5 @@ def main(argv: Sequence[str]) -> None:
 
 
 if __name__ == "__main__":
-  app.run(main)
+  with transformer_engine_context():
+    app.run(main)

src/maxdiffusion/models/attention_flax.py

@@ -1011,7 +1011,7 @@ def __init__(
         ),
     )
 
-    self.drop_out = nnx.Dropout(dropout)
+    self.drop_out = nnx.Dropout(dropout, deterministic=False)
 
     self.norm_q = nnx.data(None)
     self.norm_k = nnx.data(None)

src/maxdiffusion/models/ltx2/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2026 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""