@@ -213,9 +213,7 @@ def run(config):
 
   for prompt_idx, current_prompt in enumerate(prompts):
     prompt_word_count = len(current_prompt.split())
-    enhance_prompt = (
-        prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold
-    )
+    enhance_prompt = prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold
 
     s0 = time.perf_counter()
     images = pipeline(
 
@@ -503,9 +503,13 @@ def device_put_replicated(x, sharding):
 
   to also shard an array based on sharding.
   """
+  if isinstance(x, jax.Array) and hasattr(x, "sharding") and x.sharding == sharding:
+    return x
+
   arr = getattr(x, "value", x)
+  arr_np = np.asarray(arr)
   shd = getattr(sharding, "value", sharding)
-  res = jax.make_array_from_callback(arr.shape, shd, lambda index: arr[index])
+  res = jax.make_array_from_callback(arr_np.shape, shd, lambda index: arr_np[index])
   if hasattr(x, "set_value"):
     x.set_value(res)
     return x
 
@@ -21,6 +21,7 @@
 import pathlib
 import shutil
 import subprocess
+import tempfile
 import unittest
 import pytest
 from absl.testing import absltest
@@ -75,7 +76,8 @@ def setUp(self):
   def test_make_dreambooth_train_iterator(self):
     instance_class_gcs_dir = "gs://maxdiffusion-github-runner-test-assets/datasets/dreambooth/instance_class"
     class_class_gcs_dir = "gs://maxdiffusion-github-runner-test-assets/datasets/dreambooth/class_class"
-    local_dir = "/tmp/"
+    self.temp_dir = tempfile.TemporaryDirectory()
+    local_dir = self.temp_dir.name
     instance_class_local_dir = max_utils.download_blobs(instance_class_gcs_dir, local_dir)
     class_class_local_dir = max_utils.download_blobs(class_class_gcs_dir, local_dir)
 
@@ -136,6 +138,7 @@ def test_make_dreambooth_train_iterator(self):
 
     cleanup(instance_class_local_dir)
     cleanup(class_class_local_dir)
+    self.temp_dir.cleanup()
 
   @pytest.mark.skip(
       "This test is deprecated and will be removed in a future version. Reason: stable diffusion 2 base is no longer in HuggingFace"
 
@@ -20,6 +20,7 @@
 import os
 import pprint
 import threading
+from typing import Any
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 from flax.training import train_state
@@ -49,6 +50,54 @@ def _to_array(x):
   return x
 
 
+def _normalize_path(path: tuple[Any, ...]) -> tuple[Any, ...]:
+  """Normalizes a PyTree path tuple for optimizer state matching."""
+  normalized = []
+  for p in path:
+    if hasattr(p, "key"):
+      val = p.key
+    elif hasattr(p, "idx"):
+      val = p.idx
+    elif hasattr(p, "name"):
+      val = p.name
+    else:
+      val = p
+
+    if isinstance(val, str) and val.isdigit():
+      val = int(val)
+    normalized.append(val)
+
+  if normalized and normalized[-1] == "value":
+    normalized = normalized[:-1]
+
+  return tuple(normalized)
+
+
+def restore_optimizer_state_by_path(
+    target_opt_shardings: Any,
+    restored_opt_state: Any,
+) -> Any:
+  """Restores optimizer state by matching exact parameter paths and applying sharding."""
+  raw_restored_map, _ = jax.tree_util.tree_flatten_with_path(restored_opt_state)
+
+  normalized_restored_map = {}
+  for raw_path, val in raw_restored_map:
+    norm_path = _normalize_path(raw_path)
+    normalized_restored_map[norm_path] = val
+
+  def _restore_leaf(path, target_sharding):
+    norm_path = _normalize_path(path)
+    if norm_path not in normalized_restored_map:
+      raise KeyError(f"Optimizer checkpoint missing path: {norm_path}.")
+    val = normalized_restored_map[norm_path]
+    return max_utils.device_put_replicated(val, target_sharding)
+
+  return jax.tree_util.tree_map_with_path(
+      _restore_leaf,
+      target_opt_shardings,
+  )
+
+
 def generate_sample(config, pipeline, filename_prefix):
   """
   Generates a video to validate training did not corrupt the model
@@ -178,9 +227,11 @@ def start_training(self):
     with nn_partitioning.axis_rules(self.config.logical_axis_rules):
       pipeline, opt_state, step = self.checkpointer.load_checkpoint()
     restore_args = {}
-    if opt_state and step:
-      restore_args = {"opt_state": opt_state, "step": step}
+    if opt_state is not None:
+      restore_args["opt_state"] = opt_state
       del opt_state
+    if step is not None:
+      restore_args["step"] = step
     if self.config.enable_ssim:
       # Generate a sample before training to compare against generated sample after training.
       pretrained_video_path = generate_sample(self.config, pipeline, filename_prefix="pre-training-")
@@ -256,26 +307,76 @@ def eval(self, mesh, eval_rng_key, step, p_eval_step, state, scheduler_state, wr
   def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data_iterator, restore_args: dict = {}):
     mesh = pipeline.mesh
     graphdef, params, rest_of_state = nnx.split(pipeline.transformer, nnx.Param, ...)
+    start_step = 0
 
     with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       state = TrainState.create(
           apply_fn=graphdef.apply, params=params, tx=optimizer, graphdef=graphdef, rest_of_state=rest_of_state
       )
-      if restore_args:
-        step = restore_args.get("step", 0)
-        max_logging.log(f"Restoring optimizer and resuming from step {step}")
-        state.replace(opt_state=restore_args.get("opt_state"), step=restore_args.get("step", 0))
-        del restore_args["opt_state"]
-        del optimizer
       state = jax.tree.map(_to_array, state)
       state_spec = nnx.get_partition_spec(state)
       state = jax.lax.with_sharding_constraint(state, state_spec)
       state_shardings = nnx.get_named_sharding(state, mesh)
-      if jax.process_index() == 0 and restore_args:
-        max_logging.log("--- Optimizer State Sharding Spec (opt_state) ---")
-        pretty_string = pprint.pformat(state_spec.opt_state, indent=4, width=60)
-        max_logging.log(pretty_string)
-        max_logging.log("------------------------------------------------")
+      if restore_args.get("opt_state") is not None:
+        max_logging.log("Restoring optimizer")
+        resharded_opt_state = restore_optimizer_state_by_path(
+            target_opt_shardings=state_shardings.opt_state,
+            restored_opt_state=restore_args["opt_state"],
+        )
+        state = state.replace(opt_state=resharded_opt_state)
+        del restore_args["opt_state"]
+        if jax.process_index() == 0:
+          max_logging.log("--- Optimizer State Sharding Spec (opt_state) ---")
+          pretty_string = pprint.pformat(state_spec.opt_state, indent=4, width=60)
+          max_logging.log(pretty_string)
+          max_logging.log("------------------------------------------------")
+
+        if restore_args.get("step") is not None:
+          step = restore_args["step"]
+          # A checkpoint labelled N is written *after* iteration N completes,
+          # and the loop starts at step 0, so it already reflects N+1 applied
+          # updates. Resume at N+1 so that iteration N is not re-run, and so
+          # that TrainState.step stays aligned with the optimizer's internal
+          # step count, which actually drives the LR schedule.
+          state = state.replace(step=step + 1)
+          start_step = step + 1
+          max_logging.log(f"Resuming training from step {start_step}")
+      else:
+        # Weights-only: no optimizer state means no schedule position.
+        # Resetting to 0 keeps the applied LR and the logged LR consistent.
+        if "step" in restore_args:
+          max_logging.log(
+              "Weights-only checkpoint detected (no opt_state); resetting step"
+              " to 0. The LR schedule will restart from the beginning. If you"
+              " intended to resume training, re-save with optimizer state."
+          )
+          weights_step = restore_args.pop("step")
+
+          # The loop restarts at 0 but checkpoint_dir still holds steps up to
+          # weights_step, and Orbax will skip saving the checkpoints at those
+          # steps. Therefore, we redirect the three output paths that
+          # pyconfig derives from run_name, then rebuild the CheckpointManager.
+          previous_run_name = self.config.run_name
+          previous_checkpoint_dir = self.config.checkpoint_dir
+          new_run_name = f"{previous_run_name}_from_weights_{weights_step}"
+          keys = self.config.get_keys()
+          keys["run_name"] = new_run_name
+          for key, subdir in (
+              ("tensorboard_dir", "tensorboard"),
+              ("checkpoint_dir", "checkpoints"),
+          ):
+            keys[key] = os.path.join(self.config.output_dir, new_run_name, subdir, "")
+          max_logging.log(
+              f"Redirecting outputs from run_name {previous_run_name} to"
+              f" {new_run_name}: {previous_checkpoint_dir} already holds"
+              f" checkpoints up to step {weights_step}, which this run would"
+              " otherwise fail to save over."
+          )
+          max_logging.log(f"  checkpoint_dir:  {self.config.checkpoint_dir}")
+          max_logging.log(f"  tensorboard_dir: {self.config.tensorboard_dir}")
+          self.checkpointer.checkpoint_manager.close()
+          self.checkpointer = self._get_checkpointer()
+
     if self.config.hardware != "gpu":
       max_utils.delete_pytree(params)
     data_shardings = self.get_data_shardings(mesh)
@@ -301,8 +402,6 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
 
     rng = jax.random.key(self.config.seed)
     rng, eval_rng_key = jax.random.split(rng)
-    start_step = 0
-    last_step_completion = datetime.datetime.now()
     local_metrics_file = open(self.config.metrics_file, "a", encoding="utf8") if self.config.metrics_file else None
     running_gcs_metrics = [] if self.config.gcs_metrics else None
     first_profiling_step = self.config.skip_first_n_steps_for_profiler
@@ -311,9 +410,6 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
     last_profiling_step = np.clip(
         first_profiling_step + self.config.profiler_steps - 1, first_profiling_step, self.config.max_train_steps - 1
     )
-    if restore_args.get("step", 0):
-      max_logging.log(f"Resuming training from step {step}")
-    start_step = restore_args.get("step", 0)
     per_device_tflops, _, _ = BaseWanTrainer.calculate_tflops(pipeline)
     scheduler_state = pipeline.scheduler_state
     example_batch = load_next_batch(train_data_iterator, None, self.config)
@@ -367,7 +463,10 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         writer.flush()
       if self.config.save_final_checkpoint:
         max_logging.log(f"Saving final checkpoint for step {step}")
-        self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state.params)
+        if self.config.save_optimizer:
+          self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state)
+        else:
+          self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state.params)
         self.checkpointer.checkpoint_manager.wait_until_finished()
       # load new state for trained transformer
       pipeline.transformer = nnx.merge(state.graphdef, state.params, state.rest_of_state)