lightx2v_train big refactor (#1325)

Author: helloyongyang

lightx2v_train/cache_data.py

@@ -0,0 +1,66 @@
+import argparse
+
+import torch
+from loguru import logger
+
+from lightx2v_train.data import build_data, build_sample_processor
+from lightx2v_train.model_zoo import build_model
+from lightx2v_train.runtime import (
+    cleanup_distributed,
+    init_distributed,
+    load_config,
+    setup_logger,
+)
+from lightx2v_train.trainers import build_trainer
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Build a training dataset cache with LightX2V.")
+    parser.add_argument("--config", required=True, help="Path to a YAML config file.")
+    parser.add_argument("--output_dir", required=True, help="Directory for cache_data.jsonl and cached tensors.")
+    parser.add_argument("--save_dtype", choices=("bf16", "fp16", "fp32"), default="bf16")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--overwrite", action="store_true")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    config = load_config(args.config)
+    config["training_cache"] = {
+        "output_dir": args.output_dir,
+        "save_dtype": args.save_dtype,
+        "seed": args.seed,
+        "overwrite": args.overwrite,
+    }
+    config["data"]["train"]["preserve_records"] = True
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    init_distributed(config)
+    setup_logger(config)
+
+    try:
+        sample_processor = build_sample_processor(config)
+        dataloader = build_data(config, train_or_val="train", sample_processor=sample_processor)
+
+        model = build_model(config)
+        model.load_components(
+            load_transformer=False,
+            load_vae=True,
+            load_condition_encoder=True,
+        )
+
+        trainer = build_trainer(config)
+        trainer.set_model(model)
+        trainer.set_data(dataloader)
+        trainer.train()
+    except Exception:
+        logger.exception("Training cache build failed")
+        raise
+    finally:
+        cleanup_distributed()
+
+
+if __name__ == "__main__":
+    main()

lightx2v_train/configs/infer/lingbot_video_t2v_dmd_lora.yaml

@@ -34,7 +34,6 @@ data:
         name: prompt_dataset
         data_path: /data/nvme5/gushiqiao/datatets/lingbot_video/vidprom_filtered_extended_first5k_auto_negative.jsonl
         max_samples: 5
-        batch_size: 1
         num_workers: 0
         pin_memory: true
         shuffle: false

lightx2v_train/configs/infer/longcat_image_edit_lora.yaml

@@ -0,0 +1,50 @@
+model:
+    name: longcat_image_edit
+    pretrained_model_name_or_path: /path/to/LongCat-Image-Edit
+    attention_backend: _native_flash
+    running_dtype: bf16
+
+data:
+    processor:
+        target_area: 1048576 # 1024 * 1024
+    val:
+        name: image_dataset
+        num_workers: 4
+        shuffle: false
+        data_path:
+            - /path/to/LightX2V_train_data_examples/dataset_edit/val.jsonl
+
+scheduler:
+    num_train_timesteps: 1000
+    timestep_distribution: logitnormal
+    logitnormal_mean: 0.0
+    logitnormal_std: 1.0
+    min_t: 0.001
+    max_t: 1.0
+    time_shift_settings:
+        do_time_shift: true
+        shift_type: exponential
+        time_shift_power: 1.0
+        dynamic_shift: true
+        patch_size: [2, 2]
+        shift_x1: 256
+        shift_x2: 4096
+        shift_y1: 0.5
+        shift_y2: 1.15
+
+inference:
+    method: image_infer
+    negative_prompt: ""
+    default_width: 1024
+    default_height: 1024
+    num_inference_steps: 50
+    enable_cfg: true
+    cfg_guidance_scale: 4.5
+    seed: 42
+    output_dir: ./output_infer/longcat_image_edit_lora
+    lora_config:
+        path: /path/to/pytorch_lora_weights.safetensors
+
+logging:
+    rank_zero_only: true
+    infer_log_every_steps: 10

lightx2v_train/configs/infer/wan.yaml

@@ -37,7 +37,6 @@ data:
         data_path:
             - /path/to/prompts/wan_t2v_prompts.txt
         max_samples: 1
-        batch_size: 1
         num_workers: 0
         shuffle: false
 

lightx2v_train/configs/infer/wan2_1_t2v_14b_tf_chunkwise_ar.yaml

@@ -37,7 +37,6 @@ data:
         data_path:
             - /path/to/wan_t2v_prompts.txt
         max_samples: 1
-        batch_size: 1
         num_workers: 0
         shuffle: false