Add WAN-VACE training functionality

Introduces training support for WAN-VACE models.

New files:
- train_wan_vace.py: Main training script.
- wan_vace_trainer.py: Trainer class for WAN-VACE.
- wan_vace_checkpointing_2_1.py: Checkpointing logic for WAN-VACE.

Co-authored-by: martinarroyo <martinarroyo@google.com>

Author: ninatu

src/maxdiffusion/checkpointing/wan_vace_checkpointer_2_1.py

@@ -0,0 +1,112 @@
+"""Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import json
+from typing import Optional, Tuple
+import jax
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
+from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
+import numpy as np
+import orbax.checkpoint as ocp
+from .. import max_logging
+from ..pipelines.wan.wan_vace_pipeline_2_1 import VaceWanPipeline2_1
+
+
+class WanVaceCheckpointer2_1(WanCheckpointer):
+
+  def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dict], Optional[int]]:
+    if step is None:
+      step = self.checkpoint_manager.latest_step()
+      max_logging.log(f"Latest WAN checkpoint step: {step}")
+      if step is None:
+        max_logging.log("No WAN checkpoint found.")
+        return None, None
+    max_logging.log(f"Loading WAN checkpoint from step {step}")
+
+    cpu_devices = np.array(jax.devices(backend="cpu"))
+    mesh = Mesh(cpu_devices, axis_names=("data",))
+    replicated_sharding = NamedSharding(mesh, P())
+
+    metadatas = self.checkpoint_manager.item_metadata(step)
+    state = metadatas.wan_state
+
+    def add_sharding_to_struct(leaf_struct, sharding):
+      return jax.ShapeDtypeStruct(
+          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
+      )
+
+    target_shardings = jax.tree_util.tree_map(
+        lambda x: replicated_sharding, state
+    )
+
+    with mesh:
+      abstract_train_state_with_sharding = jax.tree_util.tree_map(
+          add_sharding_to_struct, state, target_shardings
+      )
+
+    max_logging.log("Restoring WAN checkpoint")
+    restored_checkpoint = self.checkpoint_manager.restore(
+        step=step,
+        args=ocp.args.Composite(
+            wan_config=ocp.args.JsonRestore(),
+            wan_state=ocp.args.StandardRestore(
+                abstract_train_state_with_sharding
+            ),
+        ),
+    )
+    max_logging.log(f"restored checkpoint {restored_checkpoint.keys()}")
+    max_logging.log(f"restored checkpoint wan_state {restored_checkpoint.wan_state.keys()}")
+    max_logging.log(f"optimizer found in checkpoint {'opt_state' in restored_checkpoint.wan_state.keys()}")
+    max_logging.log(f"optimizer state saved in attribute self.opt_state {self.opt_state}")
+    return restored_checkpoint, step
+
+  def load_diffusers_checkpoint(self):
+    pipeline = VaceWanPipeline2_1.from_pretrained(self.config)
+    return pipeline
+
+  def load_checkpoint(self, step=None) -> Tuple[VaceWanPipeline2_1, Optional[dict], Optional[int]]:
+    restored_checkpoint, step = self.load_wan_configs_from_orbax(step)
+    opt_state = None
+    if restored_checkpoint:
+      max_logging.log("Loading WAN pipeline from checkpoint")
+      pipeline = VaceWanPipeline2_1.from_checkpoint(self.config, restored_checkpoint)
+      if "opt_state" in restored_checkpoint.wan_state.keys():
+        opt_state = restored_checkpoint.wan_state["opt_state"]
+    else:
+      max_logging.log("No checkpoint found, loading default pipeline.")
+      pipeline = self.load_diffusers_checkpoint()
+
+    return pipeline, opt_state, step
+
+  def save_checkpoint(
+      self, train_step, pipeline: VaceWanPipeline2_1, train_states: dict
+  ):
+    """Saves the training state and model configurations."""
+
+    def config_to_json(model_or_config):
+      return json.loads(model_or_config.to_json_string())
+
+    max_logging.log(f"Saving checkpoint for step {train_step}")
+
+    # Save the checkpoint
+    self.checkpoint_manager.save(
+        train_step,
+        args=ocp.args.Composite(
+            wan_config=ocp.args.JsonSave(config_to_json(pipeline.transformer)),
+            wan_state=ocp.args.StandardSave(train_states),
+        ),
+    )
+
+    max_logging.log(f"Checkpoint for step {train_step} is saved.")

src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py

@@ -338,7 +338,14 @@ def load_transformer(
     return wan_transformer
 
   @classmethod
-  def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transformer=True):
+  def _load_and_init(
+      cls,
+      config: HyperParameters,
+      restored_checkpoint=None,
+      vae_only=False,
+      load_transformer=True,
+      load_common_components=True,
+  ):
     devices_array = max_utils.create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
     rng = jax.random.key(config.seed)
@@ -348,20 +355,31 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transform
     scheduler = None
     scheduler_state = None
     text_encoder = None
+    wan_vae = None
+    vae_cache = None
+
     if not vae_only:
       if load_transformer:
         with mesh:
           transformer = cls.load_transformer(
-              devices_array=devices_array, mesh=mesh, rngs=rngs, config=config, subfolder="transformer"
+              devices_array=devices_array,
+              mesh=mesh,
+              rngs=rngs,
+              config=config,
+              restored_checkpoint=restored_checkpoint,
+              subfolder="transformer",
           )
+      if load_common_components:
+        text_encoder = cls.load_text_encoder(config=config)
+        tokenizer = cls.load_tokenizer(config=config)
 
-      text_encoder = cls.load_text_encoder(config=config)
-      tokenizer = cls.load_tokenizer(config=config)
-
-      scheduler, scheduler_state = cls.load_scheduler(config=config)
+        scheduler, scheduler_state = cls.load_scheduler(config=config)
 
-    with mesh:
-      wan_vae, vae_cache = cls.load_vae(devices_array=devices_array, mesh=mesh, rngs=rngs, config=config)
+    if load_common_components:
+      with mesh:
+        wan_vae, vae_cache = cls.load_vae(
+            devices_array=devices_array, mesh=mesh, rngs=rngs, config=config
+        )
 
     pipeline = cls(
         tokenizer=tokenizer,
@@ -376,7 +394,43 @@ def from_pretrained(cls, config: HyperParameters, vae_only=False, load_transform
         config=config,
     )
 
-    pipeline.transformer = cls.quantize_transformer(config, pipeline.transformer, pipeline, mesh)
+    return pipeline
+
+  @classmethod
+  def from_pretrained(
+      cls,
+      config: HyperParameters,
+      vae_only=False,
+      load_transformer=True,
+      load_common_components=True,
+  ):
+    pipeline = cls._load_and_init(
+        config, None, vae_only, load_transformer, load_common_components
+    )
+    pipeline.transformer = cls.quantize_transformer(
+        config, pipeline.transformer, pipeline, pipeline.mesh
+    )
+    return pipeline
+
+  @classmethod
+  def from_checkpoint(
+      cls,
+      config: HyperParameters,
+      restored_checkpoint=None,
+      vae_only=False,
+      load_transformer=True,
+      load_common_components=True,
+  ):
+    pipeline = cls._load_and_init(
+        config,
+        restored_checkpoint,
+        vae_only,
+        load_transformer,
+        load_common_components,
+    )
+    pipeline.transformer = cls.quantize_transformer(
+        config, pipeline.transformer, pipeline, pipeline.mesh
+    )
     return pipeline
 
   def check_inputs(

src/maxdiffusion/train_wan_vace.py

@@ -0,0 +1,46 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Sequence
+
+import jax
+from absl import app
+from maxdiffusion import max_logging, pyconfig
+from maxdiffusion.train_utils import validate_train_config
+import flax
+
+
+def train(config):
+  from maxdiffusion.trainers.wan_vace_trainer import WanVaceTrainer
+
+  trainer = WanVaceTrainer(config)
+  trainer.start_training()
+
+
+def main(argv: Sequence[str]) -> None:
+  pyconfig.initialize(argv, validate_training=True)
+  config = pyconfig.config
+  validate_train_config(config)
+  max_logging.log(f"Found {jax.device_count()} devices.")
+  try:
+    flax.config.update("flax_always_shard_variable", False)
+  except LookupError:
+    pass
+  train(config)
+
+
+if __name__ == "__main__":
+  app.run(main)