Merge branch 'main' into ltx2-attention

Author: Perseus14

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/wan/transformers/transformer_wan.py

@@ -237,7 +237,7 @@ def __init__(
     else:
       raise NotImplementedError(f"{activation_fn} is not implemented.")
 
-    self.drop_out = nnx.Dropout(dropout)
+    self.drop_out = nnx.Dropout(dropout, deterministic=False)
     self.proj_out = nnx.Linear(
         rngs=rngs,
         in_features=inner_dim,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -17,7 +17,6 @@
 from typing import List, Union, Optional
 from ...pyconfig import HyperParameters
 from functools import partial
-from contextlib import nullcontext
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 import jax
@@ -116,15 +115,8 @@ def __call__(
         scheduler=self.scheduler,
         scheduler_state=scheduler_state,
     )
-    # Set the TE shard_guard context_manager if using TE cudnn_flash attention
-    if self.config.attention == "cudnn_flash_te":
-      from transformer_engine.jax.sharding import global_shard_guard, MeshResource  # pytype: disable=import-error
 
-      shard_guard = global_shard_guard(MeshResource(cp_resource="context"))
-    else:
-      shard_guard = nullcontext()
-
-    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules), shard_guard:
+    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       latents = p_run_inference(
           graphdef=graphdef,
           sharded_state=state,

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -17,7 +17,6 @@
 from typing import List, Union, Optional
 from ...pyconfig import HyperParameters
 from functools import partial
-from contextlib import nullcontext
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 import jax
@@ -140,15 +139,8 @@ def __call__(
         scheduler=self.scheduler,
         scheduler_state=scheduler_state,
     )
-    # Set the TE shard_guard context_manager if using TE cudnn_flash attention
-    if self.config.attention == "cudnn_flash_te":
-      from transformer_engine.jax.sharding import global_shard_guard, MeshResource  # pytype: disable=import-error
 
-      shard_guard = global_shard_guard(MeshResource(cp_resource="context"))
-    else:
-      shard_guard = nullcontext()
-
-    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules), shard_guard:
+    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       latents = p_run_inference(
           low_noise_graphdef=low_noise_graphdef,
           low_noise_state=low_noise_state,