Shard video_condition to prevent OOM in WAN 2.2 I2V

src/maxdiffusion/models/ltx_video/transformers_pytorch/attention.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import inspect
 from importlib import import_module
 from typing import Any, Dict, Optional, Tuple

src/maxdiffusion/pedagogical_examples/attention_comparison.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import time
 

src/maxdiffusion/pedagogical_examples/dataset_tf_cache_to_tfrecord.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import argparse
 import tensorflow as tf

src/maxdiffusion/pedagogical_examples/parameter_count.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import Sequence
 from absl import app
 import jax

src/maxdiffusion/pipelines/controlnet/__init__.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import TYPE_CHECKING
 
 from ...utils import (

src/maxdiffusion/pipelines/flux/__init__.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 _import_structure = {"pipeline_jflux": "JfluxPipeline"}
 
 from .flux_pipeline import (

src/maxdiffusion/pipelines/stable_diffusion/__init__.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import TYPE_CHECKING
 
 from ...utils import (

src/maxdiffusion/pipelines/stable_diffusion_xl/__init__.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import TYPE_CHECKING
 
 from ...utils import (

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -543,8 +543,9 @@ def prepare_latents_i2v_base(
 
     vae_dtype = getattr(self.vae, "dtype", jnp.float32)
     video_condition = video_condition.astype(vae_dtype)
-
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      sharding_spec = P(self.config.mesh_axes[0], None, None, None, None)
+      video_condition = jax.lax.with_sharding_constraint(video_condition, sharding_spec)
       encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode()
 
     # Normalize latents

src/maxdiffusion/tests/configuration_utils_test.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import json
 import os
 

src/maxdiffusion/tests/flop_calculations_test.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import unittest
 from unittest.mock import Mock

src/maxdiffusion/tests/generate_flux_smoke_test.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import unittest
 import pytest

src/maxdiffusion/tests/generate_sdxl_smoke_test.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import unittest
 import pytest

src/maxdiffusion/tests/generate_smoke_test.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 import unittest
 import pytest

src/maxdiffusion/utils/deprecation_utils.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import inspect
 import warnings
 from typing import Any, Dict, Optional, Union

src/maxdiffusion/utils/export_utils.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import io
 import random
 import struct

src/maxdiffusion/utils/loading_utils.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import os
 from typing import Callable, List, Optional, Union
 

src/maxdiffusion/utils/pil_utils.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from typing import List
 
 import PIL.Image

src/maxdiffusion/utils/testing_utils.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import functools
 import importlib
 import inspect