Shard video_condition to prevent OOM in WAN 2.2 I2V (#313)

* sharding before call to vae encoder

* sharding before call to vae encoder

* ruff check

* pyink checks

* pyink check

Author: prishajain1

code_style.sh

@@ -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/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 os
 import time
 

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 argparse
 import tensorflow as tf

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.
 """
+
 from typing import Sequence
 from absl import app
 import jax

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 TYPE_CHECKING
 
 from ...utils import (

src/maxdiffusion/pipelines/controlnet/__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/flux/__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/__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

@@ -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