feat(ltx2): pass sharding specs to VAE and embeddings connector

Author: Perseus14

src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple, Union, Optional, Sequence
+from typing import Tuple, Union, Optional, Sequence, Any
 
 import jax
 import jax.numpy as jnp
@@ -584,6 +584,7 @@ def __init__(
       dtype: jnp.dtype = jnp.float32,
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
+      sharding_specs: Optional[Any] = None,
   ):
     if timestep_conditioning:
       self.time_embedder = nnx.data(
@@ -594,6 +595,7 @@ def __init__(
               use_additional_conditions=False,
               dtype=dtype,
               weights_dtype=weights_dtype,
+              sharding_specs=sharding_specs,
           )
       )
     else:
@@ -674,6 +676,7 @@ def __init__(
       dtype: jnp.dtype = jnp.float32,
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
+      sharding_specs: Optional[Any] = None,
   ):
     out_channels = out_channels or in_channels
 
@@ -687,6 +690,7 @@ def __init__(
               use_additional_conditions=False,
               dtype=dtype,
               weights_dtype=weights_dtype,
+              sharding_specs=sharding_specs,
           )
       )
 
@@ -960,6 +964,7 @@ def __init__(
       dtype: jnp.dtype = jnp.float32,
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
+      sharding_specs: Optional[Any] = None,
   ):
     self.patch_size = patch_size
     self.patch_size_t = patch_size_t
@@ -999,6 +1004,7 @@ def __init__(
         dtype=dtype,
         weights_dtype=weights_dtype,
         precision=precision,
+        sharding_specs=sharding_specs,
     )
 
     # up blocks
@@ -1026,6 +1032,7 @@ def __init__(
               dtype=dtype,
               weights_dtype=weights_dtype,
               precision=precision,
+              sharding_specs=sharding_specs,
           )
       )
 
@@ -1059,6 +1066,7 @@ def __init__(
               use_additional_conditions=False,
               dtype=dtype,
               weights_dtype=weights_dtype,
+              sharding_specs=sharding_specs,
           )
       )
     else:
@@ -1155,6 +1163,7 @@ def __init__(
       dtype: jnp.dtype = jnp.float32,
       weights_dtype: jnp.dtype = jnp.float32,
       precision: jax.lax.Precision = None,
+      sharding_specs: Optional[Any] = None,
   ):
     self.encoder = LTX2VideoEncoder3d(
         in_channels=in_channels,
@@ -1196,6 +1205,7 @@ def __init__(
         dtype=dtype,
         weights_dtype=weights_dtype,
         precision=precision,
+        sharding_specs=sharding_specs,
     )
 
     self.scaling_factor = scaling_factor

src/maxdiffusion/models/ltx2/text_encoders/embeddings_connector_ltx2.py

@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Any
 import jax
 import jax.numpy as jnp
 from flax import nnx
@@ -37,6 +37,7 @@ def __init__(
       attention_kernel: str = "flash",
       mesh: jax.sharding.Mesh = None,
       rngs: nnx.Rngs = None,
+      sharding_specs: Optional[Any] = None,
   ):
     self.attn1 = LTX2Attention(
         query_dim=dim,
@@ -48,8 +49,9 @@ def __init__(
         attention_kernel=attention_kernel,
         mesh=mesh,
         rngs=rngs,
+        sharding_specs=sharding_specs,
     )
-    self.ff = NNXSimpleFeedForward(rngs=rngs, dim=dim, dim_out=dim, activation_fn="gelu_tanh")
+    self.ff = NNXSimpleFeedForward(rngs=rngs, dim=dim, dim_out=dim, activation_fn="gelu_tanh", sharding_specs=sharding_specs)
     self.norm1 = nnx.RMSNorm(dim, epsilon=1e-6, dtype=jnp.float32, param_dtype=jnp.float32, use_scale=False, rngs=rngs)
     self.norm2 = nnx.RMSNorm(dim, epsilon=1e-6, dtype=jnp.float32, param_dtype=jnp.float32, use_scale=False, rngs=rngs)
 
@@ -92,6 +94,7 @@ def __init__(
       attention_kernel: str = "flash",
       mesh: jax.sharding.Mesh = None,
       rngs: nnx.Rngs = None,
+      sharding_specs: Optional[Any] = None,
   ):
     self.dim = input_dim
     self.heads = heads
@@ -117,6 +120,7 @@ def create_block(rngs):
           attention_kernel=attention_kernel,
           mesh=mesh,
           rngs=rngs,
+          sharding_specs=sharding_specs,
       )
 
     # Call the vmapped constructor