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

Author: Perseus14

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -329,13 +329,18 @@ def load_connectors(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, co
     max_logging.log("Loading Connectors...")
 
     def create_model(rngs: nnx.Rngs, config: HyperParameters):
+      sharding_config = getattr(config, "sharding", {})
+      connector_strategy = sharding_config.get("text_connector", "default")
+      connector_specs = get_sharding_specs(connector_strategy, "text_connector")
+
       connectors = LTX2AudioVideoGemmaTextEncoder.from_config(
           config.pretrained_model_name_or_path,
           subfolder="connectors",
           rngs=rngs,
           mesh=mesh,
           dtype=jnp.float32,
           weights_dtype=config.weights_dtype if hasattr(config, "weights_dtype") else jnp.float32,
+          sharding_specs=connector_specs,
       )
       return connectors
 
@@ -371,13 +376,18 @@ def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: H
     max_logging.log("Loading Video VAE...")
 
     def create_model(rngs: nnx.Rngs, config: HyperParameters):
+      sharding_config = getattr(config, "sharding", {})
+      vae_strategy = sharding_config.get("vae", "default")
+      vae_specs = get_sharding_specs(vae_strategy, "vae")
+
       vae = LTX2VideoAutoencoderKL.from_config(
           config.pretrained_model_name_or_path,
           subfolder="vae",
           rngs=rngs,
           mesh=mesh,
           dtype=jnp.float32,
           weights_dtype=config.weights_dtype if hasattr(config, "weights_dtype") else jnp.float32,
+          sharding_specs=vae_specs,
       )
       return vae