changed ckpt name

Author: Serenagu525

requirements.txt

@@ -23,7 +23,6 @@ pytest==8.2.2
 tensorflow>=2.17.0
 tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
-git+https://github.com/mlperf/logging.git
 git+https://github.com/Lightricks/LTX-Video
 git+https://github.com/zmelumian972/xla@torchax/jittable_module_callable#subdirectory=torchax
 opencv-python-headless==4.10.0.84

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -213,7 +213,7 @@ def load_state_if_possible(
     max_logging.log(f"restoring from this run's directory latest step {latest_step}")
     try:
       if not enable_single_replica_ckpt_restoring:
-        if checkpoint_item == " ":
+        if checkpoint_item == "ltxvid_transformer":
           return checkpoint_manager.restore(latest_step, args=ocp.args.StandardRestore(abstract_unboxed_pre_state))
         else:
           item = {checkpoint_item: orbax.checkpoint.args.PyTreeRestore(item=abstract_unboxed_pre_state)}

src/maxdiffusion/generate_ltx_video.py

@@ -19,7 +19,7 @@
 from typing import Sequence
 from maxdiffusion.pipelines.ltx_video.ltx_video_pipeline import LTXVideoPipeline
 from maxdiffusion.pipelines.ltx_video.ltx_video_pipeline import LTXMultiScalePipeline
-from maxdiffusion import pyconfig
+from maxdiffusion import pyconfig, max_logging
 import imageio
 from datetime import datetime
 import os
@@ -108,7 +108,7 @@ def run(config):
       enhance_prompt=enhance_prompt,
       seed=config.seed,
   )
-  print("generation time: ", (time.perf_counter() - s0))
+  max_logging.log(f"Compile time: {time.perf_counter() - s0:.1f}s.")
 
   (pad_left, pad_right, pad_top, pad_bottom) = padding
   pad_bottom = -pad_bottom
@@ -146,7 +146,6 @@ def run(config):
           resolution=(height, width, config.num_frames),
           dir=output_dir,
       )
-      print(output_filename)
       # Write video
       with imageio.get_writer(output_filename, fps=fps) as video:
         for frame in video_np:

src/maxdiffusion/max_utils.py

@@ -405,7 +405,7 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
-        if checkpoint_item == " ":
+        if checkpoint_item == "ltxvid_transformer":
           state = state
         else:
           state = state[checkpoint_item]

src/maxdiffusion/models/ltx_video/linear.py

@@ -95,8 +95,6 @@ def compute_dot_general(inputs, kernel, axis, contract_ind):
     axis = _normalize_axes(axis, inputs.ndim)
 
     kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features
-    # kernel_in_axis = np.arange(len(axis))
-    # kernel_out_axis = np.arange(len(axis), len(axis) + len(features))
     kernel = self.param(
         "kernel",
         nn.with_logical_partitioning(self.kernel_init, self.kernel_axes),

src/maxdiffusion/models/ltx_video/transformers/adaln.py

@@ -126,7 +126,6 @@ def __call__(self, timesteps: jnp.ndarray) -> jnp.ndarray:
 
 
 class AlphaCombinedTimestepSizeEmbeddings(nn.Module):
-  """ """
 
   embedding_dim: int
   size_emb_dim: int