changed input format

Author: Serenagu525

src/maxdiffusion/models/ltx_video/utils/conversion_script_instruction.md

@@ -1,10 +1,3 @@
 ### Transformer Pytorch Weight Downloading and Jax Weight Loading Instructions:
-1. Weight Downloading and Conversion
-    - If first time running (no local safetensors): \
-    In the src/maxdiffusion/models/ltx_video/utils folder, run python convert_torch_weights_to_jax.py --download_ckpt_path [location to download safetensors] --output_dir [location to save jax ckpt] --transformer_config_path ../xora_v1.2-13B-balanced-128.json.
-    - If already have local pytorch checkpoint: \
-    Replace the --download_ckpt_path with --local_ckpt_path and add corresponding location
-2. Restoring Jax Weights into transformer:
-    - Replace the "ckpt_path" in src/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json with jax ckpt path.
-    - Run python src/maxdiffusion/generate_ltx_video.py src/maxdiffusion/configs/ltx_video.yml in the outer repo folder.
-
+In the folder src/maxdiffusion/models/ltx_video/utils, run:
+python convert_torch_weights_to_jax.py --ckpt_path [LOCAL DIRECTORY FOR WEIGHTS] --transformer_config_path ../xora_v1.2-13B-balanced-128.json

src/maxdiffusion/models/ltx_video/utils/convert_torch_weights_to_jax.py

@@ -230,18 +230,16 @@ def main(args):
   weight_file = "ltxv-13b-0.9.7-dev.safetensors"
 
   # download from huggingface, otherwise load from local
-  if args.local_ckpt_path is None:
-    print("Loading from HF", flush=True)
-    model_name = "Lightricks/LTX-Video"
-    local_file_path = hf_hub_download(
-        repo_id=model_name,
-        filename=weight_file,
-        local_dir=args.download_ckpt_path,
-        local_dir_use_symlinks=False,
-    )
-  else:
-    base_dir = args.local_ckpt_path
-    local_file_path = os.path.join(base_dir, weight_file)
+
+  print("Loading from HF", flush=True)
+  model_name = "Lightricks/LTX-Video"
+  absolute_ckpt_path = os.path.abspath(args.ckpt_path)
+  local_file_path = hf_hub_download(
+      repo_id=model_name,
+      filename=weight_file,
+      local_dir=absolute_ckpt_path,
+      local_dir_use_symlinks=False,
+  )
   torch_state_dict = load_file(local_file_path)
 
   print("Initializing pytorch transformer..", flush=True)
@@ -284,7 +282,7 @@ def main(args):
   params_jax = torch_statedict_to_jax(params_jax, torch_state_dict)
 
   print("Creating checkpointer and jax state for saving..", flush=True)
-  relative_ckpt_path = args.output_dir
+  relative_ckpt_path = os.path.join(args.ckpt_path, "jax_weights")
   absolute_ckpt_path = os.path.abspath(relative_ckpt_path)
   tx = optax.adamw(learning_rate=1e-5)
   with jax.default_device("cpu"):
@@ -303,25 +301,12 @@ def main(args):
 if __name__ == "__main__":
   parser = argparse.ArgumentParser(description="Convert Torch checkpoints to Jax format.")
   parser.add_argument(
-      "--local_ckpt_path",
+      "--ckpt_path",
       type=str,
       required=False,
       help="Local path of the checkpoint to convert. If not provided, will download from huggingface for example '/mnt/ckpt/00536000' or '/opt/dmd-torch-model/ema.pt'",
   )
 
-  parser.add_argument(
-      "--download_ckpt_path",
-      type=str,
-      required=False,
-      help="Location to download safetensors from huggingface",
-  )
-
-  parser.add_argument(
-      "--output_dir",
-      type=str,
-      required=True,
-      help="Path to save the checkpoint to. for example 'gs://lt-research-mm-europe-west4/jax_trainings/converted-from-torch'",
-  )
   parser.add_argument(
       "--output_step_num",
       default=1,