diff --git a/lightx2v_train/cache_data.py b/lightx2v_train/cache_data.py new file mode 100644 index 000000000..6cddca83c --- /dev/null +++ b/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() diff --git a/lightx2v_train/configs/infer/lingbot_video_t2v_dmd_lora.yaml b/lightx2v_train/configs/infer/lingbot_video_t2v_dmd_lora.yaml index 49ca778e9..f4f64c7be 100644 --- a/lightx2v_train/configs/infer/lingbot_video_t2v_dmd_lora.yaml +++ b/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 diff --git a/lightx2v_train/configs/infer/longcat_image_edit_lora.yaml b/lightx2v_train/configs/infer/longcat_image_edit_lora.yaml new file mode 100644 index 000000000..588b98408 --- /dev/null +++ b/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 diff --git a/lightx2v_train/configs/infer/wan.yaml b/lightx2v_train/configs/infer/wan.yaml deleted file mode 100644 index a2cc480e0..000000000 --- a/lightx2v_train/configs/infer/wan.yaml +++ /dev/null @@ -1,75 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: bf16 - vae_dtype: fp32 - t5_dtype: bf16 - load_vae: true - load_text_encoder: true - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - -distributed: - backend: nccl - timeout_minutes: 60 - fsdp2: - enabled: true - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - pin_memory: true - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -inference: - method: wan_ti2v_5b_infer - default_width: 1280 - default_height: 704 - num_frames: 121 - num_inference_steps: 50 - fps: 24 - enable_cfg: true - cfg_guidance_scale: 5.0 - seed: 42 - video_quality: 6.0 - macro_block_size: 16 - output_dir: ./output_infer/wan2_2_ti2v_5b - -logging: - rank_zero_only: true - infer_log_every_steps: 10 diff --git a/lightx2v_train/configs/infer/wan2_1_t2v_14b_tf_chunkwise_ar.yaml b/lightx2v_train/configs/infer/wan2_1_t2v_14b_tf_chunkwise_ar.yaml index b7658edc5..34c857f4b 100644 --- a/lightx2v_train/configs/infer/wan2_1_t2v_14b_tf_chunkwise_ar.yaml +++ b/lightx2v_train/configs/infer/wan2_1_t2v_14b_tf_chunkwise_ar.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 diff --git a/lightx2v_train/configs/infer/wan2_1_t2v_1_3b_tf_chunkwise_ar.yaml b/lightx2v_train/configs/infer/wan2_1_t2v_1_3b_tf_chunkwise_ar.yaml index 4c26f1ada..aed57c1fa 100644 --- a/lightx2v_train/configs/infer/wan2_1_t2v_1_3b_tf_chunkwise_ar.yaml +++ b/lightx2v_train/configs/infer/wan2_1_t2v_1_3b_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 diff --git a/lightx2v_train/configs/infer/wan2_2_ti2v_5b_tf_ar.yaml b/lightx2v_train/configs/infer/wan2_2_ti2v_5b_tf_ar.yaml deleted file mode 100644 index ef31546f9..000000000 --- a/lightx2v_train/configs/infer/wan2_2_ti2v_5b_tf_ar.yaml +++ /dev/null @@ -1,87 +0,0 @@ -model: - name: wan_ti2v_5b_ar - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - #checkpoint_path: /path/to/LightX2V/lightx2v_train/output_train/wan2_2_ti2v_5b_tf_cache_96f_sp4/checkpoint-000005000 - checkpoint_strict: true - latent_channels: 48 - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: bf16 - vae_dtype: fp32 - t5_dtype: bf16 - t5_cpu: true - load_vae: true - load_text_encoder: true - load_transformer: true - num_frame_per_chunk: 8 - local_attn_size: 32 - sink_size: 8 - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: true - size: 4 - fsdp2: - enabled: true - size: 2 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 10 - batch_size: 1 - num_workers: 0 - pin_memory: true - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -inference: - method: wan_ti2v_5b_ar_infer - default_width: 1280 - default_height: 704 - num_frames: 381 - num_inference_steps: 50 - fps: 24 - enable_cfg: true - cfg_guidance_scale: 3.0 - seed: 42 - save_latents_only: false - video_quality: 6.0 - macro_block_size: 16 - output_dir: ./output_infer/wan2_2_ti2v_5b_tf_ar_1 - -logging: - rank_zero_only: true - infer_log_every_steps: 10 diff --git a/lightx2v_train/configs/train/consistency/qwen_image_cm_cd_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_cm_cd_lora.yaml new file mode 100644 index 000000000..80346a48b --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_cm_cd_lora.yaml @@ -0,0 +1,109 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + # Optional overrides let CD distill a different Qwen-Image checkpoint. + # teacher: + # pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.0 + target_area: 1048576 # 1024 * 1024 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl + val: + name: image_dataset + num_workers: 8 + shuffle: false + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/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 + +training: + method: consistency + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + consistency: + algorithm: cm + mode: cd + time_pair: + mapping: ect + q: 2.0 + ratio_limit: 0.999 + kimg_per_stage: 1.0 + min_r: 0.0 + safety_epsilon: 0.000001 + loss: + distance: pseudo_huber + huber_constant: 0.00000001 + weighting: inverse_delta + normalize_by_numel: false + computation_dtype: float32 + teacher: + guidance_scale: 4.0 + negative_prompt: " " + cfg_norm: none + lora: + rank: 16 + alpha: 16 + target_modules: + - to_k + - to_q + - to_v + - to_out.0 + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_cm_cd_lora + +inference: + method: image_infer + negative_prompt: " " + default_width: 1024 + default_height: 1024 + num_inference_steps: 4 + enable_cfg: true + cfg_guidance_scale: 4.0 + seed: 42 + output_dir: ./output_infer/qwen_image_cm_cd_lora + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/consistency/qwen_image_cm_ct_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_cm_ct_lora.yaml new file mode 100644 index 000000000..c22126ad8 --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_cm_ct_lora.yaml @@ -0,0 +1,104 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + target_area: 1048576 # 1024 * 1024 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl + val: + name: image_dataset + num_workers: 8 + shuffle: false + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/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 + +training: + method: consistency + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + consistency: + algorithm: cm + mode: ct + time_pair: + # FastGen/ECT sigmoid mapping. With global batch 1, one stage is + # 1000 optimizer steps; scale this with the actual global batch. + mapping: ect + q: 2.0 + ratio_limit: 0.999 + kimg_per_stage: 1.0 + min_r: 0.0 + safety_epsilon: 0.000001 + loss: + distance: pseudo_huber + huber_constant: 0.00000001 + weighting: inverse_delta + normalize_by_numel: false + computation_dtype: float32 + lora: + rank: 16 + alpha: 16 + target_modules: + - to_k + - to_q + - to_v + - to_out.0 + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_cm_ct_lora + +inference: + method: image_infer + negative_prompt: " " + default_width: 1024 + default_height: 1024 + num_inference_steps: 4 + enable_cfg: true + cfg_guidance_scale: 4.0 + seed: 42 + output_dir: ./output_infer/qwen_image_cm_ct_lora + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/consistency/qwen_image_mean_flow_ct_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_mean_flow_ct_lora.yaml new file mode 100644 index 000000000..cb269ed9d --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_mean_flow_ct_lora.yaml @@ -0,0 +1,83 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + target_area: 1048576 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl + +scheduler: + num_train_timesteps: 1000 + timestep_distribution: logitnormal + logitnormal_mean: -1.1 + logitnormal_std: 2.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 + +training: + method: consistency + train_type: lora + max_train_iters: 10000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + consistency: + algorithm: mean_flow + mode: ct + sampling: + # The remaining samples use r=t and retain a flow-matching anchor. + random_endpoint_probability: 0.5 + jvp: + method: finite_difference + epsilon: 0.0001 + loss: + type: opt_grad + norm_method: poly_1.0 + norm_constant: 0.1 + tangent_warmup_steps: 0 + spatially_normalized_tangent: false + lora: + rank: 16 + alpha: 16 + target_modules: [to_k, to_q, to_v, to_out.0] + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.0 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_mean_flow_ct_lora + +inference: + # MeanFlow ODE sampling must pass r=t_next on every step; the standard + # Qwen pipeline cannot express that endpoint yet. + method: none + +logging: + rank_zero_only: true + train_log_every_iters: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml new file mode 100644 index 000000000..1683e3a6e --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml @@ -0,0 +1,103 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + # Optional override for the frozen RF teacher. + # teacher: + # pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.0 + target_area: 1048576 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl + val: + name: image_dataset + num_workers: 8 + shuffle: false + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/val.jsonl + +scheduler: + num_train_timesteps: 1000 + timestep_distribution: uniform # PCM uses its own discrete solver grid. + 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 + +training: + method: consistency + train_type: lora + max_train_iters: 20000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 500 + save_total_limit: 10 + consistency: + algorithm: pcm + mode: cd + solver: + num_solver_steps: 100 + num_phases: 4 + # Omit boundary_time to use scheduler.min_t, matching the released + # PCM code's first non-zero RF sigma. Set 0.0 for an exact x0 edge. + loss: + distance: pseudo_huber + huber_constant: 0.001 + computation_dtype: float32 + teacher: + # The released SD3 script uses cond + 3*(cond-uncond), equivalent + # to standard CFG scale 4 used by this framework. + guidance_scale: 4.0 + negative_prompt: " " + cfg_norm: none + lora: + rank: 16 + alpha: 16 + target_modules: [to_k, to_q, to_v, to_out.0] + optimizer: + learning_rate: 0.000005 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.001 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_pcm_cd_lora + +inference: + method: image_infer + negative_prompt: " " + default_width: 1024 + default_height: 1024 + num_inference_steps: ${training.consistency.solver.num_phases} + pcm_solver_steps: ${training.consistency.solver.num_solver_steps} + # PCM already distills guided teacher trajectories; extra CFG is optional. + enable_cfg: false + cfg_guidance_scale: 1.0 + seed: 42 + output_dir: ./output_infer/qwen_image_pcm_cd_lora + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 1 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/consistency/qwen_image_scm_ct_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_scm_ct_lora.yaml new file mode 100644 index 000000000..cbd555fa1 --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_scm_ct_lora.yaml @@ -0,0 +1,81 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + target_area: 1048576 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.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 + +training: + method: consistency + train_type: lora + max_train_iters: 10000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + consistency: + algorithm: scm + mode: ct + sigma_data: 0.5 + jvp: + # Finite differences are much more memory-friendly for Qwen-Image. + # Set to exact to use torch.func.jvp with math attention. + method: finite_difference + epsilon: 0.001 + loss: + tangent_warmup_steps: 10000 + tangent_warmup_constant: 0.1 + prior_weighting: true + spatially_normalized_tangent: true + normalize_by_numel: true + lora: + rank: 16 + alpha: 16 + target_modules: [to_k, to_q, to_v, to_out.0] + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.0 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_scm_ct_lora + +inference: + method: none + +logging: + rank_zero_only: true + train_log_every_iters: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/consistency/qwen_image_tcm_ct_lora.yaml b/lightx2v_train/configs/train/consistency/qwen_image_tcm_ct_lora.yaml new file mode 100644 index 000000000..3dce7f59e --- /dev/null +++ b/lightx2v_train/configs/train/consistency/qwen_image_tcm_ct_lora.yaml @@ -0,0 +1,93 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + # Optional model overrides for the frozen stage-1 CM. + # stage1: + # pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + target_area: 1048576 + shuffle: true + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl + +scheduler: + num_train_timesteps: 1000 + timestep_distribution: logitnormal + logitnormal_mean: 0.0 + logitnormal_std: 1.0 + min_t: 0.5 + 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 + +training: + method: consistency + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + consistency: + algorithm: tcm + mode: ct + # Must point to the completed stage-1 CM checkpoint. Both the frozen + # stage-1 branch and the trainable stage-2 branch initialize from it. + stage1_checkpoint: ./output_train/qwen_image_cm_ct_lora/checkpoint-000003000 + transition_time: 0.5 + boundary_probability: 0.25 + boundary_weight: 0.1 + time_pair: + mapping: ect + q: 2.0 + ratio_limit: 0.999 + kimg_per_stage: 1.0 + min_r: 0.0 + safety_epsilon: 0.000001 + loss: + distance: pseudo_huber + huber_constant: 0.00000001 + weighting: inverse_delta + normalize_by_numel: false + computation_dtype: float32 + lora: + rank: 16 + alpha: 16 + target_modules: [to_k, to_q, to_v, to_out.0] + optimizer: + learning_rate: 0.00005 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.0 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_tcm_ct_lora + +inference: + # TCM inference needs the composite stage-1/stage-2 denoiser; the native + # Qwen pipeline only owns one transformer, so periodic inference is off. + method: none + +logging: + rank_zero_only: true + train_log_every_iters: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml deleted file mode 100644 index 3280370a1..000000000 --- a/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml +++ /dev/null @@ -1,156 +0,0 @@ -model: - name: flux2_dev - pretrained_model_name_or_path: /path/to/FLUX.2-dev - max_sequence_length: 512 - text_encoder_out_layers: [10, 20, 30] - # FLUX.2-dev uses embedded guidance in the transformer, not negative-prompt CFG. - guidance_scale: 4.0 - running_dtype: bf16 - -distributed: - backend: nccl - timeout_minutes: 120 - fsdp2: - enabled: true - reshard_after_forward: - root_reshard: true - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: image_dataset - num_workers: 8 - prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 - shuffle: true - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl - val: - name: image_dataset - num_workers: 8 - shuffle: false - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/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 - shift_mu_strategy: flux2_empirical - shift_mu_num_steps: 50 - # Flux2 dev latents are already 2x2-patchified before scheduler shift length calculation. - patch_size: [1, 1] # [H, W] - -training: - method: dmd - max_train_iters: 1000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 10 - dmd: - num_inference_steps: 4 - fake_update_ratio: 2 - # Start at 512px for memory; raise to 1024px only after validating available GPUs. - image_sizes: - - [512, 512] - renoise_sigma_min: 0.02 - renoise_sigma_max: 1.0 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - random_schedule: - enabled: true - num_steps_min: 2 - num_steps_max: 16 - sigma_min: 0.02 - sigma_max: 0.98 - sampling_method: stratified - cdm: - enabled: true - norm_clip_min: 0.1 - weight: 1.0 - warmup_iters: 200 - student: - train_type: lora - lora: - rank: 16 - alpha: 16 - target_modules: - - to_q - - to_k - - to_v - - to_out.0 - - add_q_proj - - add_k_proj - - add_v_proj - - to_add_out - - to_qkv_mlp_proj - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - fake: - train_type: lora - lora: - rank: 16 - alpha: 16 - target_modules: - - to_q - - to_k - - to_v - - to_out.0 - - add_q_proj - - add_k_proj - - add_v_proj - - to_add_out - - to_qkv_mlp_proj - optimizer: - learning_rate: 0.00002 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - teacher: - # Manual negative-prompt CFG is disabled; embedded guidance comes from model.guidance_scale. - guidance_scale: 1.0 - negative_prompt: "" - cfg_norm: none - output_dir: ./output_train/flux2_dev_cdm_lora - -inference: - method: image_infer - negative_prompt: "" - default_width: 1024 - default_height: 1024 - num_inference_steps: 4 - enable_cfg: false - cfg_guidance_scale: ${model.guidance_scale} - seed: 42 - output_dir: ./output_infer/flux2_dev_cdm_lora - infer_every_iters: - -logging: - rank_zero_only: true - train_log_every_iters: 1 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/flux2_dev_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/flux2_dev_dmd_lora.yaml index d9b51e144..890c9c4bf 100644 --- a/lightx2v_train/configs/train/dmd/flux2_dev_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/flux2_dev_dmd_lora.yaml @@ -23,13 +23,14 @@ distributed: data: train: - name: image_dataset + # JSONL row: {"prompt": "...", "target_height": 512, "target_width": 512} + name: prompt_dataset num_workers: 8 + bucket_by_size: true prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 shuffle: true data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl + - /path/to/datasets/image_dmd/train.jsonl val: name: image_dataset num_workers: 8 @@ -69,7 +70,7 @@ training: fake_update_ratio: 2 # Start at 512px for memory; raise to 1024px only after validating available GPUs. image_sizes: - - [512, 512] + - {value: [512, 512]} renoise_sigma_min: 0.02 renoise_sigma_max: 1.0 renoise_discrete_samples: 1000 diff --git a/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml index a48ffccf6..507f3e4eb 100644 --- a/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml @@ -1,18 +1,19 @@ model: name: flux2_klein - pretrained_model_name_or_path: /path/to/FLUX.2-klein-4B + pretrained_model_name_or_path: /path/to/FLUX.2-klein-base-4B max_sequence_length: 1024 running_dtype: bf16 data: train: - name: image_dataset + # JSONL row: {"prompt": "...", "target_height": 1024, "target_width": 1024} + name: prompt_dataset num_workers: 8 + bucket_by_size: true prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 shuffle: true data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl + - /path/to/datasets/image_dmd/train.jsonl val: name: image_dataset num_workers: 8 @@ -43,9 +44,9 @@ training: num_inference_steps: 4 fake_update_ratio: 2 image_sizes: - - [1024, 1024] - - [768, 1344] - - [1344, 768] + - {value: [1024, 1024], ratio: 5} + - {value: [16, 768, 1344], ratio: 3} + - {value: [16, 1344, 768], ratio: 2} renoise_sigma_min: 0.02 renoise_sigma_max: 1.0 renoise_discrete_samples: 1000 diff --git a/lightx2v_train/configs/train/dmd/lingbot_video_t2v_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/lingbot_video_t2v_dmd_lora.yaml index d116f029c..e10430e46 100644 --- a/lightx2v_train/configs/train/dmd/lingbot_video_t2v_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/lingbot_video_t2v_dmd_lora.yaml @@ -39,7 +39,6 @@ data: train: name: prompt_dataset data_path: /data/nvme5/gushiqiao/datatets/lingbot_video/vidprom_filtered_extended_first5k_auto_negative.jsonl - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -48,7 +47,6 @@ data: name: prompt_dataset data_path: /data/nvme5/gushiqiao/codes/new/LightX2V/lightx2v_train/assets/lingbot_video/t2v_prompts_first5.json max_samples: 1 - batch_size: 1 num_workers: 0 pin_memory: true shuffle: false @@ -68,7 +66,7 @@ scheduler: extra_one_step: true training: - method: lingbot_video_dmd + method: dmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/lingbot_video_t2v_droid_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/lingbot_video_t2v_droid_dmd_lora.yaml index cd7fdf284..9cb024426 100644 --- a/lightx2v_train/configs/train/dmd/lingbot_video_t2v_droid_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/lingbot_video_t2v_droid_dmd_lora.yaml @@ -39,7 +39,6 @@ data: train: name: prompt_dataset data_path: /data/nvme5/gushiqiao/datatets/droid_prompts/droid_prompts_50k_lingbot_auto_negative.jsonl - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -48,7 +47,6 @@ data: name: prompt_dataset data_path: /data/nvme5/gushiqiao/codes/new/LightX2V/lightx2v_train/assets/lingbot_video/t2v_prompts_first5.json max_samples: 2 - batch_size: 1 num_workers: 0 pin_memory: true shuffle: false @@ -68,7 +66,7 @@ scheduler: extra_one_step: true training: - method: lingbot_video_dmd + method: dmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml deleted file mode 100644 index 2b800df1e..000000000 --- a/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml +++ /dev/null @@ -1,131 +0,0 @@ -model: - name: longcat_image - pretrained_model_name_or_path: /path/to/LongCat-Image - max_sequence_length: 1024 - attention_backend: _native_flash - running_dtype: bf16 - -data: - train: - name: image_dataset - num_workers: 8 - prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 - shuffle: true - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl - val: - name: image_dataset - num_workers: 8 - shuffle: false - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/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 - -training: - method: dmd - max_train_iters: 1000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 10 - dmd: - num_inference_steps: 4 - fake_update_ratio: 2 - image_sizes: - - [1024, 1024] - - [768, 1344] - - [1344, 768] - renoise_sigma_min: 0.02 - renoise_sigma_max: 1.0 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - random_schedule: - enabled: true - num_steps_min: 2 - num_steps_max: 16 - sigma_min: 0.02 - sigma_max: 0.98 - sampling_method: stratified - cdm: - enabled: true - norm_clip_min: 0.1 - weight: 1.0 - warmup_iters: 200 - student: - train_type: lora - lora: - rank: 16 - alpha: 16 - target_modules: - - to_k - - to_q - - to_v - - to_out.0 - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - fake: - train_type: lora - lora: - rank: 16 - alpha: 16 - target_modules: - - to_k - - to_q - - to_v - - to_out.0 - optimizer: - learning_rate: 0.00002 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 4.5 - negative_prompt: " " - cfg_norm: layer_norm - output_dir: ./output_train/longcat_image_cdm_lora - -inference: - method: image_infer - default_width: 1024 - default_height: 1024 - num_inference_steps: 4 - negative_prompt: " " - enable_cfg: false - cfg_guidance_scale: 1.0 - enable_prompt_rewrite: false - output_dir: ./output_infer/longcat_image_cdm_lora - infer_every_iters: - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml index 6a1b65d8e..d21acb06a 100644 --- a/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml @@ -1,19 +1,19 @@ model: name: longcat_image pretrained_model_name_or_path: /path/to/LongCat-Image - max_sequence_length: 1024 attention_backend: _native_flash running_dtype: bf16 data: train: - name: image_dataset + # JSONL row: {"prompt": "...", "target_height": 1024, "target_width": 1024} + name: prompt_dataset num_workers: 8 + bucket_by_size: true prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 shuffle: true data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl + - /path/to/datasets/image_dmd/train.jsonl val: name: image_dataset num_workers: 8 @@ -53,9 +53,9 @@ training: num_inference_steps: 4 fake_update_ratio: 2 image_sizes: - - [1024, 1024] - - [768, 1344] - - [1344, 768] + - {value: [1024, 1024], ratio: 5} + - {value: [16, 768, 1344], ratio: 3} + - {value: [16, 1344, 768], ratio: 2} renoise_sigma_min: 0.02 renoise_sigma_max: 1.0 renoise_discrete_samples: 1000 diff --git a/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml deleted file mode 100644 index 77b9aecf5..000000000 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml +++ /dev/null @@ -1,128 +0,0 @@ -model: - name: ltx_t2av_ar - pretrained_model_name_or_path: /data/nvme4/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - text_encoder_path: /data/nvme0/gushiqiao/models/official_models/LTX-2/ - running_dtype: bf16 - transformer_param_dtype: bf16 - embeddings_dtype: bf16 - text_encoder_dtype: bf16 - text_encoder_cpu: false - load_transformer: true - load_embeddings_processor: true - load_text_encoder: false - num_frame_per_chunk: 3 - local_attn_size: -1 - sink_size: 0 - defer_kv_cache_updates: false - detach_kv_cache_updates: true - fake: - name: ltx_t2av - pretrained_model_name_or_path: /data/nvme4/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - transformer_param_dtype: bf16 - teacher: - name: ltx_t2av - pretrained_model_name_or_path: /data/nvme4/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - transformer_param_dtype: bf16 - student: - checkpoint_path: /data/nvme5/gushiqiao/codes/LightX2V/lightx2v_train/output_train/ltx_t2av_teacher_forcing/checkpoint-000000900 - checkpoint_strict: true - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /data/nvme5/gushiqiao/datatets/x2v_online/lightx2v_ltx_t2av/prompt_cache/metadata.jsonl - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true - drop_last: true - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: ltx_t2av_ar_dmd - max_train_iters: 1000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 2 - ltx2: - default_fps: 25 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - dmd: - num_inference_steps: 4 - fake_update_ratio: 5 - num_frame_per_chunk: 3 - context_noise: 0.0 - height: 768 - width: 512 - num_frames: 113 - fps: 25 - renoise_sigma_min: 0.02 - renoise_sigma_max: 0.98 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - student: - train_type: full - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: full - optimizer: - learning_rate: 0.00002 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: ./output_train/ltx_t2av_ar_dmd_lora - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml deleted file mode 100644 index 878fdae33..000000000 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml +++ /dev/null @@ -1,176 +0,0 @@ -model: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - text_encoder_path: /path/to/models/LTX-2 - running_dtype: bf16 - transformer_param_dtype: bf16 - embeddings_dtype: bf16 - text_encoder_dtype: bf16 - text_encoder_cpu: false - load_transformer: true - load_embeddings_processor: true - load_text_encoder: false - fake: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - teacher: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - transformer_param_dtype: bf16 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/datasets/x2v_online/lightx2v_ltx_t2av/prompt_cache - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true - drop_last: true - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - -training: - method: ltx_t2av_dmd - max_train_iters: 1000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 10 - ltx2: - default_fps: 25 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - dmd: - num_inference_steps: 4 - fake_update_ratio: 5 - height: 768 - width: 512 - num_frames: 241 - fps: 25 - video_latent_shape: [128, 31, 24, 16] - audio_latent_shape: [8, 241, 16] - renoise_sigma_min: 0.02 - renoise_sigma_max: 0.98 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - student: - train_type: lora - lora: - rank: 128 - alpha: 128 - target_modules: - - attn1.to_q - - attn1.to_k - - attn1.to_v - - attn1.to_out.0 - - attn2.to_q - - attn2.to_k - - attn2.to_v - - attn2.to_out.0 - - ff.net.0.proj - - ff.net.2 - - audio_attn1.to_q - - audio_attn1.to_k - - audio_attn1.to_v - - audio_attn1.to_out.0 - - audio_attn2.to_q - - audio_attn2.to_k - - audio_attn2.to_v - - audio_attn2.to_out.0 - - audio_ff.net.0.proj - - audio_ff.net.2 - - audio_to_video_attn.to_q - - audio_to_video_attn.to_k - - audio_to_video_attn.to_v - - audio_to_video_attn.to_out.0 - - video_to_audio_attn.to_q - - video_to_audio_attn.to_k - - video_to_audio_attn.to_v - - video_to_audio_attn.to_out.0 - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: lora - lora: - rank: 128 - alpha: 128 - target_modules: - - attn1.to_q - - attn1.to_k - - attn1.to_v - - attn1.to_out.0 - - attn2.to_q - - attn2.to_k - - attn2.to_v - - attn2.to_out.0 - - ff.net.0.proj - - ff.net.2 - - audio_attn1.to_q - - audio_attn1.to_k - - audio_attn1.to_v - - audio_attn1.to_out.0 - - audio_attn2.to_q - - audio_attn2.to_k - - audio_attn2.to_v - - audio_attn2.to_out.0 - - audio_ff.net.0.proj - - audio_ff.net.2 - - audio_to_video_attn.to_q - - audio_to_video_attn.to_k - - audio_to_video_attn.to_v - - audio_to_video_attn.to_out.0 - - video_to_audio_attn.to_q - - video_to_audio_attn.to_k - - video_to_audio_attn.to_v - - video_to_audio_attn.to_out.0 - optimizer: - learning_rate: 0.00002 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: /path/to/LightX2V/lightx2v_train/output_train/ltx_t2av_dmd_lora - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/minimax_h3_t2av_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/minimax_h3_t2av_dmd_lora.yaml index bf95334cf..fffb313b1 100644 --- a/lightx2v_train/configs/train/dmd/minimax_h3_t2av_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/minimax_h3_t2av_dmd_lora.yaml @@ -1,5 +1,7 @@ model: name: minimax_h3_t2av + # This must be the converted Diffusers model root containing + # transformer/config.json (MiniMaxH3Transformer3DModel). pretrained_model_name_or_path: /data/nvme6/gushiqiao/models/MiniMax-H3 running_dtype: bf16 transformer_param_dtype: bf16 @@ -8,8 +10,22 @@ model: video_latent_channels: 24 audio_latent_channels: 32 vae_spatial_scale_factor: 16 + video_fps: 24 + audio_sampling_rate: 32000 + audio_latents_per_second: 40 + # Keep frozen cache encoders off GPU outside their forward passes. + cache_encoder_cpu_offload: true + # H3 keeps its diffusion latents in fp32; prompt embeddings remain bf16. + use_autocast: false # Full SDPA is impractical at H3 sequence lengths; this matches run_h3.py. attention_backend: _flash_3_hub + capabilities: + distribution_matching: + # The H3 denoiser applies separate shifts to one shared base sigma. + video_flow_shift: 6.0 + audio_flow_shift: 3.0 + video_loss_weight: 1.0 + audio_loss_weight: 1.0 distributed: backend: nccl @@ -32,9 +48,9 @@ distributed: data: train: - name: latent_dataset - data_path: /data/nvme6/gushiqiao/datasets/minimax_h3_prompt_cache - batch_size: 1 + name: prompt_dataset + data_path: + - /path/to/datasets/minimax_h3_t2av/prompts.txt num_workers: 2 pin_memory: true shuffle: true @@ -47,7 +63,7 @@ scheduler: max_t: 1.0 training: - method: minimax_h3_t2av_dmd + method: dmd max_train_iters: 1000 gradient_accumulation_iters: 1 gradient_checkpointing: true @@ -56,12 +72,6 @@ training: lr_warmup_iters: 0 save_every_iters: 100 save_total_limit: 5 - minimax_h3: - video_flow_shift: 6.0 - audio_flow_shift: 3.0 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - latent_dtype: fp32 dmd: # Four student model evaluations; the terminal zero sigma is added internally. num_inference_steps: 4 @@ -70,8 +80,11 @@ training: width: 1344 # 124 = 17*7+5, about 5.17 seconds at H3's fixed 24 fps. num_frames: 124 - renoise_sigma_min: 0.02 - renoise_sigma_max: 0.98 + score_sampling: + type: continuous_uniform + min_sigma: 0.02 + max_sigma: 0.98 + # Do not apply the generic scheduler shift before H3's modality shifts. warp_denoising_step: false student: train_type: lora diff --git a/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml deleted file mode 100644 index 68cca7fd1..000000000 --- a/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml +++ /dev/null @@ -1,143 +0,0 @@ -model: - name: qwen_image - pretrained_model_name_or_path: /path/to/Qwen/Qwen-Image - fake: - name: qwen_image - pretrained_model_name_or_path: /path/to/Qwen/Qwen-Image - teacher: - name: qwen_image - pretrained_model_name_or_path: /path/to/Qwen/Qwen-Image - max_sequence_length: 1024 - running_dtype: bf16 - -data: - train: - name: image_dataset - num_workers: 8 - prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 - shuffle: true - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl - val: - name: image_dataset - num_workers: 8 - shuffle: false - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/val.jsonl - -scheduler: - num_train_timesteps: 1000 - time_shift_settings: - do_time_shift: true - shift_type: linear - # shift function: "linear" => mu/(mu+(1/t-1)^p), "exponential" => exp(mu)/(exp(mu)+(1/t-1)^p) - time_shift_power: 1.0 - dynamic_shift: false - time_shift_mu: 3.0 - -training: - method: dmd - max_train_iters: 1000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 10 - dmd: - num_inference_steps: 4 - fake_update_ratio: 2 - image_sizes: - - [1024, 1024] - - [768, 1344] - - [1344, 768] - renoise_sigma_min: 0.02 - renoise_sigma_max: 1.0 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - random_schedule: - enabled: true - num_steps_min: 2 - num_steps_max: 16 - sigma_min: 0.02 - sigma_max: 0.98 - sampling_method: stratified - cdm: - enabled: true - norm_clip_min: 0.1 - weight: 1.0 - warmup_iters: 200 - student: - train_type: lora - lora: - rank: 32 - alpha: 32 - target_modules: - - to_k - - to_q - - to_v - - to_out.0 - # - add_q_proj - # - add_k_proj - # - add_v_proj - # - to_add_out - # - img_mlp.net.0.proj - # - img_mlp.net.2 - # - txt_mlp.net.0.proj - # - txt_mlp.net.2 - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - fake: - train_type: lora - lora: - rank: 32 - alpha: 32 - target_modules: - - to_k - - to_q - - to_v - - to_out.0 - # - add_q_proj - # - add_k_proj - # - add_v_proj - # - to_add_out - # - img_mlp.net.0.proj - # - img_mlp.net.2 - # - txt_mlp.net.0.proj - # - txt_mlp.net.2 - optimizer: - learning_rate: 0.00002 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.001 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 4.0 - negative_prompt: " " - cfg_norm: layer_norm - output_dir: ./output_train/qwen_image_cdm_lora - -inference: - method: image_infer - default_width: 1024 - default_height: 1024 - num_inference_steps: 4 - cfg_guidance_scale: 4.0 - negative_prompt: " " - enable_cfg: false - output_dir: ./output_infer/qwen_image_cdm_lora - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/qwen_image_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/qwen_image_dmd_lora.yaml index 5cce78929..512eedfe4 100644 --- a/lightx2v_train/configs/train/dmd/qwen_image_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/qwen_image_dmd_lora.yaml @@ -12,13 +12,14 @@ model: data: train: - name: image_dataset + # JSONL row: {"prompt": "...", "target_height": 1024, "target_width": 1024} + name: prompt_dataset num_workers: 8 + bucket_by_size: true prompt_dropout_rate: 0.0 - target_area: 1048576 # 1024 * 1024 shuffle: true data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl + - /path/to/datasets/image_dmd/train.jsonl val: name: image_dataset num_workers: 8 @@ -50,9 +51,9 @@ training: num_inference_steps: 4 fake_update_ratio: 2 image_sizes: - - [1024, 1024] - - [768, 1344] - - [1344, 768] + - {value: [1024, 1024], ratio: 5} + - {value: [16, 768, 1344], ratio: 3} + - {value: [16, 1344, 768], ratio: 2} renoise_sigma_min: 0.02 renoise_sigma_max: 1.0 renoise_discrete_samples: 1000 diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd.yaml index aaa63cdd7..4e173bed3 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd.yaml @@ -45,7 +45,6 @@ data: name: prompt_dataset data_path: - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -55,7 +54,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -73,7 +71,7 @@ scheduler: extra_one_step: true training: - method: video_ar_dmd + method: autoregressive_dmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd_lora.yaml index 8c6e07729..bf7dece75 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_ar_dmd_lora.yaml @@ -45,7 +45,6 @@ data: name: prompt_dataset data_path: - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -55,7 +54,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -73,7 +71,7 @@ scheduler: extra_one_step: true training: - method: video_ar_dmd + method: autoregressive_dmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd.yaml index 7388de254..b30a0bbc6 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd.yaml @@ -45,7 +45,6 @@ data: frame_rate: 16 fix_frame_rate: true random_start: true - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -55,7 +54,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -73,7 +71,7 @@ scheduler: extra_one_step: true training: - method: video_dmd + method: dmd max_train_iters: 5000 gradient_accumulation_iters: 1 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd_lora.yaml index bb0e97f6f..db23df2b8 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_14b_dmd_lora.yaml @@ -39,7 +39,6 @@ data: name: prompt_dataset data_path: - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -49,7 +48,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -67,7 +65,7 @@ scheduler: extra_one_step: true training: - method: video_dmd + method: dmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_ar_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_ar_dmd.yaml index 74064a4de..ead0b4b2e 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_ar_dmd.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_ar_dmd.yaml @@ -46,7 +46,6 @@ data: name: prompt_dataset data_path: - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -56,7 +55,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -74,7 +72,7 @@ scheduler: extra_one_step: true training: - method: video_ar_dmd + method: autoregressive_dmd max_train_iters: 1000 gradient_accumulation_iters: 1 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_dmd.yaml index f99a76649..1b58dd70a 100644 --- a/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_dmd.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_1_t2v_1_3b_dmd.yaml @@ -39,7 +39,6 @@ data: name: prompt_dataset data_path: - data/prompts/wan22_train_prompts.txt - batch_size: 2 num_workers: 8 pin_memory: true shuffle: true @@ -49,7 +48,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -67,9 +65,9 @@ scheduler: extra_one_step: true training: - method: video_dmd + method: dmd max_train_iters: 1000 - gradient_accumulation_iters: 1 + gradient_accumulation_iters: 2 gradient_checkpointing: true max_grad_norm: 10.0 lr_scheduler: constant diff --git a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_ar_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_ar_dmd.yaml deleted file mode 100644 index b1331a448..000000000 --- a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_ar_dmd.yaml +++ /dev/null @@ -1,175 +0,0 @@ -model: - name: wan_ti2v_5b_ar - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - latent_channels: 48 - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: fp32 - vae_dtype: fp32 - t5_dtype: bf16 - t5_cpu: false - load_vae: true - load_text_encoder: true - load_transformer: true - num_frame_per_chunk: 8 - local_attn_size: 32 - defer_kv_cache_updates: true - detach_kv_cache_updates: false - sink_size: 8 - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - fake: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - transformer_param_dtype: bf16 - teacher: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - transformer_param_dtype: bf16 - student: - checkpoint_path: /path/to/LightX2V/lightx2v_train/output_train/wan2_2_ti2v_5b_tf_cache_96f_sp4/checkpoint-000005000 - checkpoint_strict: true - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: 4 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: prompt_dataset - data_path: - - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - pin_memory: true - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: video_ar_dmd - max_train_iters: 5000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 2 - save_consolidated_student: true - dmd: - num_inference_steps: 4 - fake_update_ratio: 5 - height: 704 - width: 1280 - num_frames: 125 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - ts_schedule: false - ts_schedule_max: false - num_train_timestep: 1000 - timestep_shift: 5.0 - min_score_timestep: 0 - min_step_ratio: 0.02 - max_step_ratio: 0.98 - same_step_across_blocks: true - sp_cache: true - num_frame_per_chunk: 8 - context_noise: 0 - renoise_sigma_min: 0.02 - renoise_sigma_max: 0.98 - renoise_discrete_samples: 1000 - renoise_shift: 5.0 - student: - train_type: full - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: full - optimizer: - learning_rate: 0.0000004 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: ./output_train4/wan2_2_ti2v_5b_ar_dmd - -inference: - method: wan_ti2v_5b_ar_infer - default_width: 1280 - default_height: 704 - num_frames: 125 - num_inference_steps: 4 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - fps: 24 - enable_cfg: false - cfg_guidance_scale: 1.0 - seed: 42 - save_latents_only: false - output_dir: ./output_infer4/wan2_2_ti2v_5b_ar_dmd - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml deleted file mode 100644 index 16a021c87..000000000 --- a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml +++ /dev/null @@ -1,156 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: fp32 - vae_dtype: fp32 - t5_dtype: bf16 - load_vae: true - load_text_encoder: true - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - fake: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - teacher: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - transformer_param_dtype: bf16 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: 4 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: prompt_dataset - data_path: - - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: video_dmd - max_train_iters: 5000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 1 - dmd: - ida: - enabled: true - decay: 0.97 - num_inference_steps: 4 - fake_update_ratio: 5 - height: 704 - width: 1280 - num_frames: 121 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - ts_schedule: false - ts_schedule_max: false - num_train_timestep: 1000 - timestep_shift: 5.0 - min_score_timestep: 0 - min_step_ratio: 0.02 - max_step_ratio: 0.98 - student: - train_type: full - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: full - optimizer: - learning_rate: 0.0000004 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: ./output_train/wan2_2_ti2v_5b_dmd - -inference: - method: wan_ti2v_5b_infer - default_width: 1280 - default_height: 704 - num_frames: 121 - num_inference_steps: 4 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - fps: 24 - enable_cfg: false - cfg_guidance_scale: 1.0 - seed: 42 - output_dir: ./output_infer/wan2_2_ti2v_5b_dmd - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd_multi.yaml b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd_multi.yaml deleted file mode 100644 index 7571d4047..000000000 --- a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd_multi.yaml +++ /dev/null @@ -1,157 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: bf16 - vae_dtype: fp32 - t5_dtype: bf16 - load_vae: true - load_text_encoder: true - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - fake: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - transformer_param_dtype: bf16 - teacher: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - transformer_param_dtype: bf16 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: true - size: 2 - fsdp2: - enabled: true - size: 2 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: prompt_dataset - data_path: - - /path/to/prompts/vidprom_filtered_extended.txt - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: video_dmd - max_train_iters: 5000 - gradient_accumulation_iters: 2 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 1 - dmd: - ida: - enabled: true - decay: 0.97 - num_inference_steps: 4 - fake_update_ratio: 5 - height: 704 - width: 1280 - num_frames: 121 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - ts_schedule: false - ts_schedule_max: false - num_train_timestep: 1000 - timestep_shift: 5.0 - min_score_timestep: 0 - min_step_ratio: 0.02 - max_step_ratio: 0.98 - student: - train_type: full - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: full - optimizer: - learning_rate: 0.0000004 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: ./output_train/wan2_2_ti2v_5b_dmd_multi - -inference: - method: wan_ti2v_5b_infer - default_width: 1280 - default_height: 704 - num_frames: 121 - num_inference_steps: 4 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - fps: 24 - enable_cfg: false - cfg_guidance_scale: 1.0 - seed: 42 - output_dir: ./output_infer/wan2_2_ti2v_5b_dmd_multi - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml b/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml deleted file mode 100644 index cc0f88895..000000000 --- a/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml +++ /dev/null @@ -1,104 +0,0 @@ -model: - name: flux2_klein - pretrained_model_name_or_path: /path/to/FLUX.2-Klein-9B - text_encoder_out_layers: [9, 18, 27] - running_dtype: bf16 - -data: - train: - name: image_dataset - num_workers: 4 - prompt_dropout_rate: 0.0 - target_area: 1048576 - shuffle: true - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl - val: - name: image_dataset - num_workers: 4 - target_area: 1048576 - shuffle: false - data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/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 - shift_mu_strategy: flux2_empirical - shift_mu_num_steps: 4 - patch_size: [1, 1] - -training: - train_type: lora - method: dopsd - max_train_iters: 4000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 200 - save_total_limit: 20 - dopsd: - num_training_steps: 4 - ema_decay: 0.9999 - ema_decay_warmup: 0.999 - ema_decay_warmup_iters: 200 - # early-step loss weighting: t0, t1, t2, t3 - step_loss_weights: [8.0, 4.0, 2.0, 1.0] - # teacher edit branch only; student uses dataset prompt (T2I) - edit_sys_prompt: "The output must be exactly the same as the reference image." - # teacher prompt: false = only edit_sys_prompt (reference image carries content) - # true = "{dataset_prompt} {edit_sys_prompt}" (original D-OPSD default) - teacher_use_dataset_prompt: false - # save student/teacher x0 trajectory grid during training (not at infer) - trajectory_every_iters: ${training.save_every_iters} - lora: - rank: 64 - alpha: 128 - target_modules: - - to_q - - to_k - - to_v - - to_out.0 - - add_q_proj - - add_k_proj - - add_v_proj - - to_add_out - - to_qkv_mlp_proj - optimizer: - learning_rate: 0.00005 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.0 - adam_epsilon: 0.00000001 - output_dir: ./output_train/dopsd_editcontext_ema0.9999_onpolicy_4steptrain_9b_corgi_bsz1_lora_lr5e-5_earlyw - -inference: - method: image_infer - negative_prompt: "" - default_width: 1024 - default_height: 1024 - num_inference_steps: 4 - enable_cfg: false - cfg_guidance_scale: 1.0 - seed: 42 - output_dir: ./output_infer/dopsd_editcontext_ema0.9999_onpolicy_4steptrain_9b_corgi_bsz1_lora_lr5e-5_earlyw - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml b/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml deleted file mode 100644 index efc546dcc..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml +++ /dev/null @@ -1,67 +0,0 @@ -model: - name: wan_fastwam - running_dtype: bf16 - model_path: /path/to/Wan-AI/Wan2.2-TI2V-5B - -data: - train: &libero_plus_data - name: libero_fastwam_dataset - dataset_dirs: - # https://huggingface.co/datasets/Sylvest/libero_plus_lerobot - - /path/to/datasets/libero_plus_lerobot - shape_meta: - images: [front, wrist] - # Match the released Kairos recipe; training-time eval is in-sample. - val_set_proportion: 0.0 - batch_size: 8 - num_workers: 8 - val: - <<: *libero_plus_data - -training: - method: fastwam - output_dir: runs/fastwam_libero_plus_4gpu - # Approximately 10 epochs over 2,238,036 frames at global batch 128. - max_train_iters: 174848 - gradient_accumulation_iters: 4 - max_grad_norm: 1.0 - # One checkpoint per approximately 2,238,036 / 128 optimizer steps. - save_every_iters: 17485 - save_total_limit: 10 - save_final: true - lr_scheduler: cosine - lr_warmup_iters: 8742 - # Cosine floor = optimizer.learning_rate * lr_eta_min_ratio; set to 0 to decay to zero. - lr_eta_min_ratio: 0.01 - optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-2 - adam_beta1: 0.9 - adam_beta2: 0.95 - adam_epsilon: 1.0e-8 - -inference: - infer_every_iters: 200 - eval_num_inference_steps: 5 - num_samples: 4 - sample_strategy: rotating - run_inference: true - save_video: true - save_preview: true - fps: 8 - tiled: false - seed: 42 - -logging: - train_log_every_iters: 10 - -resume: - auto_resume: true - resume_ckpt_path: null - -distributed: - backend: nccl - zero1: - enabled: true - sequence_parallel: - enabled: false diff --git a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml b/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml deleted file mode 100644 index eb6388d8d..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml +++ /dev/null @@ -1,67 +0,0 @@ -model: - name: wan_fastwam - running_dtype: bf16 - model_path: /path/to/Wan-AI/Wan2.2-TI2V-5B - -data: - train: &libero_plus_data - name: libero_fastwam_dataset - dataset_dirs: - # https://huggingface.co/datasets/Sylvest/libero_plus_lerobot - - /path/to/datasets/libero_plus_lerobot - shape_meta: - images: [front, wrist] - # Match the released Kairos recipe; training-time eval is in-sample. - val_set_proportion: 0.0 - batch_size: 8 - num_workers: 8 - val: - <<: *libero_plus_data - -training: - method: fastwam - output_dir: runs/fastwam_libero_plus_8gpu - # Approximately 10 epochs over 2,238,036 frames at global batch 128. - max_train_iters: 174848 - gradient_accumulation_iters: 2 - max_grad_norm: 1.0 - # One checkpoint per approximately 174848 / 8 optimizer steps. - save_every_iters: 17485 - save_total_limit: 10 - save_final: true - lr_scheduler: cosine - lr_warmup_iters: 8742 - # Cosine floor = optimizer.learning_rate * lr_eta_min_ratio; set to 0 to decay to zero. - lr_eta_min_ratio: 0.01 - optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-2 - adam_beta1: 0.9 - adam_beta2: 0.95 - adam_epsilon: 1.0e-8 - -inference: - infer_every_iters: 200 - eval_num_inference_steps: 5 - num_samples: 8 - sample_strategy: rotating - run_inference: true - save_video: true - save_preview: true - fps: 8 - tiled: false - seed: 42 - -logging: - train_log_every_iters: 10 - -resume: - auto_resume: true - resume_ckpt_path: null - -distributed: - backend: nccl - zero1: - enabled: true - sequence_parallel: - enabled: false diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml deleted file mode 100644 index 4f20d265e..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ /dev/null @@ -1,64 +0,0 @@ -model: - name: wan_fastwam - running_dtype: bf16 - model_path: /path/to/Wan-AI/Wan2.2-TI2V-5B - -data: - train: &libero_data - name: libero_fastwam_dataset - dataset_dirs: - # https://huggingface.co/datasets/yuanty/LIBERO-fastwam - - /path/to/LIBERO-fastwam/libero_spatial_no_noops_lerobot - - /path/to/LIBERO-fastwam/libero_object_no_noops_lerobot - - /path/to/LIBERO-fastwam/libero_goal_no_noops_lerobot - - /path/to/LIBERO-fastwam/libero_10_no_noops_lerobot - batch_size: 8 - num_workers: 8 - val: - <<: *libero_data - -training: - method: fastwam - output_dir: runs/fastwam_libero_full - # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule. - max_train_iters: 21700 - gradient_accumulation_iters: 4 - max_grad_norm: 1.0 - save_every_iters: 2000 - save_total_limit: 2 - save_final: true - lr_scheduler: cosine - lr_warmup_iters: 1085 - # Cosine floor = optimizer.learning_rate * lr_eta_min_ratio; set to 0 to decay to zero. - lr_eta_min_ratio: 0.01 - optimizer: - learning_rate: 1.0e-4 - weight_decay: 1.0e-2 - adam_beta1: 0.9 - adam_beta2: 0.95 - adam_epsilon: 1.0e-8 - -inference: - infer_every_iters: 200 - eval_num_inference_steps: 5 - num_samples: 1 - run_inference: true - save_video: true - save_preview: true - fps: 8 - tiled: false - seed: 42 - -logging: - train_log_every_iters: 10 - -resume: - auto_resume: true - resume_ckpt_path: null - -distributed: - backend: nccl - zero1: - enabled: true - sequence_parallel: - enabled: false diff --git a/lightx2v_train/configs/train/flow/flux2_dev_edit_lora.yaml b/lightx2v_train/configs/train/flow/flux2_dev_edit_lora.yaml new file mode 100644 index 000000000..0cfb4b496 --- /dev/null +++ b/lightx2v_train/configs/train/flow/flux2_dev_edit_lora.yaml @@ -0,0 +1,106 @@ +model: + name: flux2_dev_edit + pretrained_model_name_or_path: /path/to/FLUX.2-dev + max_sequence_length: 512 + text_encoder_out_layers: [10, 20, 30] + guidance_scale: 4.0 + running_dtype: bf16 + +distributed: + backend: nccl + timeout_minutes: 120 + fsdp2: + enabled: true + reshard_after_forward: + root_reshard: false + block_reshard: true + mixed_precision: + param_dtype: bf16 + reduce_dtype: fp32 + output_dtype: null + cast_forward_inputs: false + +data: + processor: + target_area: 1048576 # 1024 * 1024 + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + shuffle: true + data_path: + - /path/to/LightX2V_train_data_examples/dataset_edit/train.jsonl + val: + name: image_dataset + num_workers: 8 + 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 + shift_mu_strategy: flux2_empirical + shift_mu_num_steps: 50 + patch_size: [1, 1] + +training: + method: flow_matching + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + lora: + rank: 16 + alpha: 16 + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - add_q_proj + - add_k_proj + - add_v_proj + - to_add_out + - to_qkv_mlp_proj + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/flux2_dev_edit_lora + +inference: + method: image_infer + negative_prompt: "" + default_width: 1024 + default_height: 1024 + num_inference_steps: 50 + enable_cfg: false + cfg_guidance_scale: ${model.guidance_scale} + seed: 42 + output_dir: ./output_infer/flux2_dev_edit_lora + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml b/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml index 7dcab5247..619de74c2 100644 --- a/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml +++ b/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml @@ -21,11 +21,12 @@ distributed: cast_forward_inputs: false data: + processor: + target_area: 1048576 # 1024 * 1024 train: name: image_dataset num_workers: 8 prompt_dropout_rate: 0.1 - target_area: 1048576 # 1024 * 1024 shuffle: true data_path: - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl @@ -54,7 +55,7 @@ scheduler: patch_size: [1, 1] # [H, W] training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 @@ -86,7 +87,6 @@ training: output_dir: ./output_train/flux2_dev_lora inference: - # Use the distributed-safe inferencer for torchrun/FSDP. Use image_native_infer only for single-process eval. method: image_infer negative_prompt: "" default_width: 1024 diff --git a/lightx2v_train/configs/train/flow/flux2_klein_edit_lora.yaml b/lightx2v_train/configs/train/flow/flux2_klein_edit_lora.yaml new file mode 100644 index 000000000..b14aa118b --- /dev/null +++ b/lightx2v_train/configs/train/flow/flux2_klein_edit_lora.yaml @@ -0,0 +1,91 @@ +model: + name: flux2_klein_edit + pretrained_model_name_or_path: /path/to/FLUX.2-klein-base-9B + max_sequence_length: 512 + text_encoder_out_layers: [9, 18, 27] + running_dtype: bf16 + +data: + processor: + target_area: 1048576 # 1024 * 1024 + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.1 + shuffle: true + data_path: + - /path/to/LightX2V_train_data_examples/dataset_edit/train.jsonl + val: + name: image_dataset + num_workers: 8 + 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 + shift_mu_strategy: flux2_empirical + shift_mu_num_steps: 50 + patch_size: [1, 1] + +training: + method: flow_matching + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 200 + save_total_limit: 10 + lora: + rank: 16 + alpha: 32 + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - add_q_proj + - add_k_proj + - add_v_proj + - to_add_out + - to_qkv_mlp_proj + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/flux2_klein_edit_lora + +inference: + method: image_infer + negative_prompt: "" + default_width: 1024 + default_height: 1024 + num_inference_steps: 50 + enable_cfg: true + cfg_guidance_scale: 4.0 + seed: 42 + output_dir: ./output_infer/flux2_klein_edit_lora + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml b/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml index 24f86bf06..72f46904b 100644 --- a/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml +++ b/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml @@ -1,16 +1,17 @@ model: name: flux2_klein - pretrained_model_name_or_path: /path/to/FLUX.2-Klein-9B + pretrained_model_name_or_path: /path/to/FLUX.2-klein-base-9B max_sequence_length: 512 text_encoder_out_layers: [9, 18, 27] running_dtype: bf16 data: + processor: + target_area: 1048576 # 1024 * 1024 train: name: image_dataset num_workers: 8 prompt_dropout_rate: 0.1 - target_area: 1048576 # 1024 * 1024 shuffle: true # examples: https://github.com/ModelTC/LightX2V_train_data_examples data_path: @@ -35,12 +36,12 @@ scheduler: time_shift_power: 1.0 dynamic_shift: true shift_mu_strategy: flux2_empirical - shift_mu_num_steps: 4 + shift_mu_num_steps: 50 # Flux2 latents are already 2x2-patchified before scheduler shift length calculation. patch_size: [1, 1] # [H, W] training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 @@ -76,9 +77,9 @@ inference: negative_prompt: "" default_width: 1024 default_height: 1024 - num_inference_steps: 4 - enable_cfg: false - cfg_guidance_scale: 1.0 + num_inference_steps: 50 + enable_cfg: true + cfg_guidance_scale: 4.0 seed: 42 output_dir: ./output_infer/flux2_klein_lora infer_every_iters: ${training.save_every_iters} diff --git a/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml b/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml index b651ba0a5..9be80d315 100644 --- a/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml +++ b/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml @@ -1,6 +1,6 @@ model: name: flux2_klein - pretrained_model_name_or_path: /path/to/FLUX.2-Klein-9B + pretrained_model_name_or_path: /path/to/FLUX.2-klein-base-9B max_sequence_length: 512 text_encoder_out_layers: [9, 18, 27] running_dtype: bf16 @@ -20,11 +20,12 @@ distributed: cast_forward_inputs: false data: + processor: + target_area: 1048576 # 1024 * 1024 train: name: image_dataset num_workers: 8 prompt_dropout_rate: 0.1 - target_area: 1048576 # 1024 * 1024 shuffle: true # examples: https://github.com/ModelTC/LightX2V_train_data_examples data_path: @@ -54,7 +55,7 @@ scheduler: patch_size: [1, 1] # [H, W] training: - method: flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/longcat_image_edit_lora.yaml b/lightx2v_train/configs/train/flow/longcat_image_edit_lora.yaml new file mode 100644 index 000000000..63a6b15ec --- /dev/null +++ b/lightx2v_train/configs/train/flow/longcat_image_edit_lora.yaml @@ -0,0 +1,87 @@ +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 + train: + name: image_dataset + num_workers: 8 + prompt_dropout_rate: 0.05 + shuffle: true + data_path: + - /path/to/LightX2V_train_data_examples/dataset_edit/train.jsonl + 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 + +training: + method: flow_matching + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + lora: + rank: 16 + alpha: 16 + target_modules: + - to_k + - to_q + - to_v + - to_out.0 + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/longcat_image_edit_lora + +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 + infer_every_iters: ${training.save_every_iters} + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/flow/longcat_image_lora.yaml b/lightx2v_train/configs/train/flow/longcat_image_lora.yaml index 2998517d0..134a0b5be 100644 --- a/lightx2v_train/configs/train/flow/longcat_image_lora.yaml +++ b/lightx2v_train/configs/train/flow/longcat_image_lora.yaml @@ -1,16 +1,16 @@ model: name: longcat_image pretrained_model_name_or_path: /path/to/LongCat-Image - max_sequence_length: 1024 attention_backend: _native_flash running_dtype: bf16 data: + processor: + target_area: 1048576 # 1024 * 1024 train: name: image_dataset num_workers: 8 prompt_dropout_rate: 0.1 - target_area: 1048576 # 1024 * 1024 shuffle: true # examples: https://github.com/ModelTC/LightX2V_train_data_examples data_path: @@ -36,14 +36,14 @@ scheduler: time_shift_power: 1.0 dynamic_shift: true patch_size: [2, 2] # [H, W] - # https://github.com/huggingface/diffusers/blob/v0.38.0/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L59 + # https://github.com/huggingface/diffusers/blob/v0.38.0/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py shift_x1: 256 shift_x2: 4096 shift_y1: 0.5 shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 @@ -76,7 +76,7 @@ inference: default_height: 1024 num_inference_steps: 50 enable_cfg: true - cfg_guidance_scale: 4.0 + cfg_guidance_scale: 4.5 seed: 42 output_dir: ./output_infer/longcat_image_lora infer_every_iters: ${training.save_every_iters} diff --git a/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml b/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml index c8ec9196c..4fa6193c3 100644 --- a/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml +++ b/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml @@ -1,7 +1,6 @@ model: name: longcat_image pretrained_model_name_or_path: /path/to/LongCat-Image - max_sequence_length: 1024 attention_backend: _native_flash running_dtype: bf16 @@ -20,11 +19,12 @@ distributed: cast_forward_inputs: false data: + processor: + target_area: 1048576 # 1024 * 1024 train: name: image_dataset num_workers: 8 prompt_dropout_rate: 0.1 - target_area: 1048576 # 1024 * 1024 shuffle: true # examples: https://github.com/ModelTC/LightX2V_train_data_examples data_path: @@ -50,14 +50,14 @@ scheduler: time_shift_power: 1.0 dynamic_shift: true patch_size: [2, 2] # [H, W] - # https://github.com/huggingface/diffusers/blob/v0.38.0/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L59 + # https://github.com/huggingface/diffusers/blob/v0.38.0/src/diffusers/pipelines/longcat_image/pipeline_longcat_image.py shift_x1: 256 shift_x2: 4096 shift_y1: 0.5 shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 1 @@ -82,7 +82,7 @@ inference: default_height: 1024 num_inference_steps: 50 enable_cfg: true - cfg_guidance_scale: 4.0 + cfg_guidance_scale: 4.5 seed: 42 output_dir: ./output_infer/longcat_image_full_fsdp2 infer_every_iters: ${training.save_every_iters} diff --git a/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml deleted file mode 100644 index e1b07768f..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml +++ /dev/null @@ -1,76 +0,0 @@ -model: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - running_dtype: bf16 - transformer_param_dtype: bf16 - embeddings_dtype: bf16 - load_transformer: true - load_embeddings_processor: true - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/datasets/lightx2v_ltx_t2av/.precomputed/metadata.jsonl - dataset_repeat: 1 - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true -scheduler: - num_train_timesteps: 1000 - timestep_distribution: shifted_logit_normal - logitnormal_std: 1.0 - logitnormal_eps: 0.001 - logitnormal_uniform_prob: 0.1 - -training: - method: ltx_t2av_flow - train_type: full - max_train_iters: 3000 - gradient_accumulation_iters: 2 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 10 - ltx2: - default_fps: 24 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - optimizer: - learning_rate: 0.000001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: ./output_train/ltx_t2av_full - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml deleted file mode 100644 index cb49f24e7..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml +++ /dev/null @@ -1,108 +0,0 @@ -model: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - running_dtype: bf16 - transformer_param_dtype: bf16 - embeddings_dtype: bf16 - load_transformer: true - load_embeddings_processor: true - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/ltx_t2av/.precomputed - dataset_repeat: 1 - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true -scheduler: - num_train_timesteps: 1000 - timestep_distribution: shifted_logit_normal - logitnormal_std: 1.0 - logitnormal_eps: 0.001 - logitnormal_uniform_prob: 0.1 - -training: - method: ltx_t2av_flow - train_type: lora - max_train_iters: 3000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 10 - ltx2: - default_fps: 24 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - lora: - rank: 64 - alpha: 64 - target_modules: - - attn1.to_q - - attn1.to_k - - attn1.to_v - - attn1.to_out.0 - - attn2.to_q - - attn2.to_k - - attn2.to_v - - attn2.to_out.0 - - ff.net.0.proj - - ff.net.2 - - audio_attn1.to_q - - audio_attn1.to_k - - audio_attn1.to_v - - audio_attn1.to_out.0 - - audio_attn2.to_q - - audio_attn2.to_k - - audio_attn2.to_v - - audio_attn2.to_out.0 - - audio_ff.net.0.proj - - audio_ff.net.2 - - audio_to_video_attn.to_q - - audio_to_video_attn.to_k - - audio_to_video_attn.to_v - - audio_to_video_attn.to_out.0 - - video_to_audio_attn.to_q - - video_to_audio_attn.to_k - - video_to_audio_attn.to_v - - video_to_audio_attn.to_out.0 - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: ./output_train/ltx_t2av_lora - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml deleted file mode 100644 index 0e95bad41..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml +++ /dev/null @@ -1,108 +0,0 @@ -model: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - running_dtype: bf16 - transformer_param_dtype: fp32 - embeddings_dtype: bf16 - load_transformer: true - load_embeddings_processor: true - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/datasets/x2v_online/lightx2v_ltx_t2av/.precomputed - dataset_repeat: 1 - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true -scheduler: - num_train_timesteps: 1000 - timestep_distribution: shifted_logit_normal - logitnormal_std: 1.0 - logitnormal_eps: 0.001 - logitnormal_uniform_prob: 0.1 - -training: - method: ltx_t2av_flow - train_type: lora - max_train_iters: 3000 - gradient_accumulation_iters: 2 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 10 - save_every_iters: 100 - save_total_limit: 6 - ltx2: - default_fps: 25 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - lora: - rank: 128 - alpha: 128 - target_modules: - - attn1.to_q - - attn1.to_k - - attn1.to_v - - attn1.to_out.0 - - attn2.to_q - - attn2.to_k - - attn2.to_v - - attn2.to_out.0 - - ff.net.0.proj - - ff.net.2 - - audio_attn1.to_q - - audio_attn1.to_k - - audio_attn1.to_v - - audio_attn1.to_out.0 - - audio_attn2.to_q - - audio_attn2.to_k - - audio_attn2.to_v - - audio_attn2.to_out.0 - - audio_ff.net.0.proj - - audio_ff.net.2 - - audio_to_video_attn.to_q - - audio_to_video_attn.to_k - - audio_to_video_attn.to_v - - audio_to_video_attn.to_out.0 - - video_to_audio_attn.to_q - - video_to_audio_attn.to_k - - video_to_audio_attn.to_v - - video_to_audio_attn.to_out.0 - optimizer: - learning_rate: 0.0001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: /path/to/LightX2V/lightx2v_train/output_train/ltx_t2av - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/flow/minimax_h3_t2av_lora.yaml b/lightx2v_train/configs/train/flow/minimax_h3_t2av_lora.yaml new file mode 100644 index 000000000..2d7489c3d --- /dev/null +++ b/lightx2v_train/configs/train/flow/minimax_h3_t2av_lora.yaml @@ -0,0 +1,116 @@ +model: + name: minimax_h3_t2av + # Converted Diffusers model root containing transformer/config.json. + pretrained_model_name_or_path: /path/to/models/MiniMax-H3 + running_dtype: bf16 + transformer_param_dtype: bf16 + # H3 diffusion states are kept in fp32; text embeddings remain bf16. + latent_dtype: fp32 + local_files_only: true + patch_size: [1, 2, 2] + video_latent_channels: 24 + audio_latent_channels: 32 + vae_spatial_scale_factor: 16 + video_fps: 24 + audio_sampling_rate: 32000 + audio_latents_per_second: 40 + # Cache construction activates the large frozen encoders one at a time. + cache_encoder_cpu_offload: true + use_autocast: false + attention_backend: _flash_3_hub + capabilities: + flow_matching: + video_flow_shift: 6.0 + audio_flow_shift: 3.0 + video_loss_weight: 1.0 + audio_loss_weight: 1.0 + +distributed: + backend: nccl + timeout_minutes: 60 + sequence_parallel: + enabled: false + size: 1 + fsdp2: + enabled: true + size: ${oc.env:WORLD_SIZE,1} + reshard_after_forward: + root_reshard: false + block_reshard: true + # Preserve the transformer's intentional fp32/bf16 parameter mix. + mixed_precision: + param_dtype: null + reduce_dtype: fp32 + output_dtype: null + cast_forward_inputs: false + +data: + train: + name: video_dataset + data_path: + - /path/to/datasets/minimax_h3_t2av/train.jsonl + video_column: video + audio_column: audio + prompt_column: prompt + height: 768 + width: 1344 + # H3 requires F = 17*n+5. + num_frames: 124 + frame_rate: 24 + fix_frame_rate: true + random_start: false + skip_missing: true + decode_retries: 3 + num_workers: 4 + pin_memory: true + shuffle: true + drop_last: true + +scheduler: + num_train_timesteps: 1000 + timestep_distribution: uniform + min_t: 0.001 + max_t: 1.0 + # The H3 capability applies separate rational shifts to video and audio. + time_shift_settings: + do_time_shift: false + +training: + method: flow_matching + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + lora: + rank: 128 + alpha: 128 + target_modules: + - to_q + - to_k + - to_v + - to_out.0 + - ff.net.0.proj + - ff.net.2 + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/minimax_h3_t2av_lora + +inference: + method: none + infer_every_iters: null + +logging: + rank_zero_only: true + train_log_every_iters: 10 + +resume: + auto_resume: false diff --git a/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora.yaml b/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora.yaml index 05ca60f19..bb83b1bc0 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora.yaml @@ -42,7 +42,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 1000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora_dist.yaml b/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora_dist.yaml index abcc8e389..4443b7b70 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora_dist.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_edit_2511_lora_dist.yaml @@ -56,7 +56,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 1000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/qwen_image_lora.yaml b/lightx2v_train/configs/train/flow/qwen_image_lora.yaml index a254124fb..1bc8e1193 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_lora.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_lora.yaml @@ -1,6 +1,6 @@ model: name: qwen_image - pretrained_model_name_or_path: /path/to/Qwen/Qwen-Image-2512 + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 max_sequence_length: 1024 running_dtype: bf16 @@ -13,13 +13,13 @@ data: shuffle: true # examples: https://github.com/ModelTC/LightX2V_train_data_examples data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/train.jsonl + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/train.jsonl val: name: image_dataset num_workers: 8 shuffle: false data_path: - - /path/to/LightX2V_train_data_examples/dataset_v1/val.jsonl + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/val.jsonl scheduler: num_train_timesteps: 1000 @@ -42,7 +42,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/qwen_image_lora_cache.yaml b/lightx2v_train/configs/train/flow/qwen_image_lora_cache.yaml new file mode 100644 index 000000000..2e43b75ba --- /dev/null +++ b/lightx2v_train/configs/train/flow/qwen_image_lora_cache.yaml @@ -0,0 +1,85 @@ +model: + name: qwen_image + pretrained_model_name_or_path: /root/yongyang/models/Qwen/Qwen-Image-2512 + max_sequence_length: 1024 + running_dtype: bf16 + +data: + use_training_cache: true + + train: + name: image_dataset + num_workers: 8 + persistent_workers: true + prefetch_factor: 2 + pin_memory: true + prompt_dropout_rate: 0.1 + target_area: 1048576 # 1024 * 1024 + shuffle: true + data_path: + - /root/yongyang/datasets/cache_data/cache_data.jsonl + val: + name: image_dataset + num_workers: 8 + shuffle: false + data_path: + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v1/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 + # shift function: "linear" => mu/(mu+(1/t-1)^p), "exponential" => exp(mu)/(exp(mu)+(1/t-1)^p) + time_shift_power: 1.0 + dynamic_shift: true + patch_size: [2, 2] # [H, W] + # https://github.com/huggingface/diffusers/blob/v0.38.0/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L59 + shift_x1: 256 + shift_x2: 4096 + shift_y1: 0.5 + shift_y2: 1.15 + +training: + method: flow_matching + train_type: lora + max_train_iters: 3000 + gradient_accumulation_iters: 1 + gradient_checkpointing: true + max_grad_norm: 1.0 + lr_scheduler: constant + lr_warmup_iters: 10 + save_every_iters: 100 + save_total_limit: 10 + lora: + rank: 16 + alpha: 16 + target_modules: + - to_k + - to_q + - to_v + - to_out.0 + optimizer: + learning_rate: 0.0001 + adam_beta1: 0.9 + adam_beta2: 0.999 + weight_decay: 0.01 + adam_epsilon: 0.00000001 + output_dir: ./output_train/qwen_image_lora + +inference: + method: none + infer_every_iters: null + +logging: + rank_zero_only: true + train_log_every_iters: 10 + infer_log_every_steps: 10 + +resume: + auto_resume: true diff --git a/lightx2v_train/configs/train/flow/qwen_image_lora_dist.yaml b/lightx2v_train/configs/train/flow/qwen_image_lora_dist.yaml index 2092f3956..c84f3a8b4 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_lora_dist.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_lora_dist.yaml @@ -60,7 +60,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/qwen_image_lora_dp.yaml b/lightx2v_train/configs/train/flow/qwen_image_lora_dp.yaml index 46e1e0d23..867dbe2b5 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_lora_dp.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_lora_dp.yaml @@ -53,7 +53,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: lora max_train_iters: 1000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/qwen_image_t2i.yaml b/lightx2v_train/configs/train/flow/qwen_image_t2i.yaml index 43b1a051d..c5d3ccf59 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_t2i.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_t2i.yaml @@ -60,7 +60,7 @@ scheduler: shift_y2: 1.15 training: - method: flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_full.yaml b/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_full.yaml index c99af27e2..49d28bccb 100644 --- a/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_full.yaml +++ b/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_full.yaml @@ -45,7 +45,6 @@ data: random_start: false skip_missing: true decode_retries: 3 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -54,7 +53,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 5 - batch_size: 1 num_workers: 0 shuffle: false @@ -78,7 +76,7 @@ scheduler: max_resolution: [720, 1280] training: - method: flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_lora.yaml b/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_lora.yaml index d2249a0a8..34122841b 100644 --- a/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_lora.yaml +++ b/lightx2v_train/configs/train/flow/wan2_1_t2v_14b_lora.yaml @@ -45,7 +45,6 @@ data: random_start: false skip_missing: true decode_retries: 3 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -54,7 +53,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 5 - batch_size: 1 num_workers: 0 shuffle: false @@ -78,7 +76,7 @@ scheduler: max_resolution: [720, 1280] training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_full.yaml b/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_full.yaml index f18ea6b23..8f95cfa46 100644 --- a/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_full.yaml +++ b/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_full.yaml @@ -45,7 +45,6 @@ data: random_start: false skip_missing: true decode_retries: 3 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -54,7 +53,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -78,7 +76,7 @@ scheduler: max_resolution: [720, 1280] training: - method: flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_lora.yaml b/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_lora.yaml index d9e10a02d..194123c2a 100644 --- a/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_lora.yaml +++ b/lightx2v_train/configs/train/flow/wan2_1_t2v_1_3b_lora.yaml @@ -1,6 +1,6 @@ model: name: wan_t2v - pretrained_model_name_or_path: /path/to/Wan2.1-T2V-1.3B + pretrained_model_name_or_path: /root/yongyang/models/Wan-AI/Wan2.1-T2V-1.3B max_sequence_length: 512 running_dtype: bf16 vae_dtype: fp32 @@ -30,10 +30,9 @@ data: train: name: video_dataset data_path: - - /path/to/OpenVid-1M/OpenVid-1M.csv - video_root: /path/to/OpenVid-1M/video - video_column: video - prompt_column: caption + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v2/train.jsonl + video_column: target_data + prompt_column: prompt height: 480 width: 832 num_frames: 81 @@ -44,16 +43,14 @@ data: random_start: false skip_missing: true decode_retries: 3 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true val: name: prompt_dataset data_path: - - /path/to/wan_t2v_prompts.txt - max_samples: 5 - batch_size: 1 + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v2/val.jsonl + max_samples: 4 num_workers: 0 shuffle: false @@ -77,7 +74,7 @@ scheduler: max_resolution: [720, 1280] training: - method: flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/phased_dmd/wan2_1_t2v_1_3b_phased_dmd.yaml b/lightx2v_train/configs/train/phased_dmd/wan2_1_t2v_1_3b_phased_dmd.yaml index 5fa0d9d2f..03bb9a3cf 100644 --- a/lightx2v_train/configs/train/phased_dmd/wan2_1_t2v_1_3b_phased_dmd.yaml +++ b/lightx2v_train/configs/train/phased_dmd/wan2_1_t2v_1_3b_phased_dmd.yaml @@ -56,7 +56,6 @@ data: frame_rate: 16 fix_frame_rate: true random_start: true - batch_size: 2 num_workers: 8 pin_memory: true shuffle: true @@ -66,7 +65,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -84,9 +82,9 @@ scheduler: extra_one_step: true training: - method: video_phased_dmd + method: phased_dmd max_train_iters: 1000 - gradient_accumulation_iters: 1 + gradient_accumulation_iters: 2 gradient_checkpointing: true max_grad_norm: 10.0 lr_scheduler: constant diff --git a/lightx2v_train/configs/train/phased_dmd/wan2_2_t2v_a14b_phased_dmd.yaml b/lightx2v_train/configs/train/phased_dmd/wan2_2_t2v_a14b_phased_dmd.yaml index 32722deb5..7c38d45f5 100644 --- a/lightx2v_train/configs/train/phased_dmd/wan2_2_t2v_a14b_phased_dmd.yaml +++ b/lightx2v_train/configs/train/phased_dmd/wan2_2_t2v_a14b_phased_dmd.yaml @@ -55,7 +55,6 @@ data: frame_rate: 16 fix_frame_rate: true random_start: true - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -65,7 +64,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -83,7 +81,7 @@ scheduler: extra_one_step: true training: - method: video_phased_dmd + method: phased_dmd max_train_iters: 4000 gradient_accumulation_iters: 1 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/sgmd/wan2_1_t2v_14b_sgmd.yaml b/lightx2v_train/configs/train/sgmd/wan2_1_t2v_14b_sgmd.yaml index d2270b4ac..c6ecf76ad 100644 --- a/lightx2v_train/configs/train/sgmd/wan2_1_t2v_14b_sgmd.yaml +++ b/lightx2v_train/configs/train/sgmd/wan2_1_t2v_14b_sgmd.yaml @@ -39,7 +39,6 @@ data: name: prompt_dataset data_path: - data/prompts/wan22_train_prompts.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -49,7 +48,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -67,7 +65,7 @@ scheduler: extra_one_step: true training: - method: video_sgmd + method: sgmd max_train_iters: 5000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/sgmd/wan2_1_t2v_1_3b_sgmd.yaml b/lightx2v_train/configs/train/sgmd/wan2_1_t2v_1_3b_sgmd.yaml index 5aef2587a..82cd979e0 100644 --- a/lightx2v_train/configs/train/sgmd/wan2_1_t2v_1_3b_sgmd.yaml +++ b/lightx2v_train/configs/train/sgmd/wan2_1_t2v_1_3b_sgmd.yaml @@ -39,7 +39,6 @@ data: name: prompt_dataset data_path: - data/prompts/wan22_train_prompts.txt - batch_size: 2 num_workers: 8 pin_memory: true shuffle: true @@ -49,7 +48,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 2 - batch_size: 1 num_workers: 0 shuffle: false @@ -67,9 +65,9 @@ scheduler: extra_one_step: true training: - method: video_sgmd + method: sgmd max_train_iters: 1000 - gradient_accumulation_iters: 1 + gradient_accumulation_iters: 2 gradient_checkpointing: true max_grad_norm: 10.0 lr_scheduler: constant diff --git a/lightx2v_train/configs/train/sgmd/wan2_2_ti2v_5b_sgmd.yaml b/lightx2v_train/configs/train/sgmd/wan2_2_ti2v_5b_sgmd.yaml deleted file mode 100644 index 309f6fe7e..000000000 --- a/lightx2v_train/configs/train/sgmd/wan2_2_ti2v_5b_sgmd.yaml +++ /dev/null @@ -1,158 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: models/Wan2.2-TI2V-5B - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: fp32 - vae_dtype: fp32 - t5_dtype: bf16 - load_vae: true - load_text_encoder: true - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - fake: - name: wan_ti2v_5b - pretrained_model_name_or_path: models/Wan2.2-TI2V-5B/ - teacher: - name: wan_ti2v_5b - pretrained_model_name_or_path: models/Wan2.2-TI2V-5B/ - transformer_param_dtype: bf16 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: 4 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: prompt_dataset - data_path: - - data/prompts/wan22_train_prompts.txt - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - data/prompts/wan22_val_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: video_sgmd - max_train_iters: 5000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 1 - dmd: - num_inference_steps: 4 - fake_update_ratio: 1 - div_loss: - enabled: false - weight: 5.0e-2 - teacher_inference_steps: 30 - anchor_step: 5 - height: 704 - width: 1280 - num_frames: 121 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - ts_schedule: false - ts_schedule_max: false - num_train_timestep: 1000 - timestep_shift: 5.0 - min_score_timestep: 0 - min_step_ratio: 0.02 - max_step_ratio: 0.98 - student: - train_type: full - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - fake: - train_type: full - optimizer: - learning_rate: 0.0000004 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - teacher: - guidance_scale: 3.0 - cfg_norm: none - output_dir: ./output_train/wan2_2_ti2v_5b_sgmd - -inference: - method: wan_ti2v_5b_infer - default_width: 1280 - default_height: 704 - num_frames: 121 - num_inference_steps: 4 - denoising_step_list: - - 1000 - - 750 - - 500 - - 250 - warp_denoising_step: true - fps: 24 - enable_cfg: false - cfg_guidance_scale: 1.0 - seed: 42 - output_dir: ./output_infer/wan2_2_ti2v_5b_sgmd - infer_every_iters: ${training.save_every_iters} - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml b/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml deleted file mode 100644 index 8513c89b4..000000000 --- a/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml +++ /dev/null @@ -1,94 +0,0 @@ -model: - name: ltx_t2av - pretrained_model_name_or_path: /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors - running_dtype: bf16 - transformer_param_dtype: bf16 - embeddings_dtype: bf16 - load_transformer: true - load_embeddings_processor: true - num_frame_per_chunk: 3 - local_attn_size: -1 - sink_size: 0 - defer_kv_cache_updates: true - detach_kv_cache_updates: false - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: false - size: 1 - fsdp2: - enabled: true - size: ${oc.env:WORLD_SIZE,1} - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/datasets/lightx2v_ltx_t2av_121/.precomputed/metadata.jsonl - dataset_repeat: 1 - batch_size: 1 - num_workers: 4 - pin_memory: true - shuffle: true - drop_last: true - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: ltx_t2av_teacher_forcing - train_type: full - max_train_iters: 3000 - gradient_accumulation_iters: 4 - gradient_checkpointing: true - max_grad_norm: 1.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 10 - save_consolidated_weights: true - teacher_forcing: - enabled: true - mode: chunkwise - num_frame_per_chunk: 3 - ltx2: - default_fps: 25 - video_loss_weight: 1.0 - audio_loss_weight: 1.0 - optimizer: - learning_rate: 0.000001 - adam_beta1: 0.9 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: ./output_train/ltx_t2av_teacher_forcing - -inference: - method: none - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - -resume: - auto_resume: false diff --git a/lightx2v_train/configs/train/tf/wan2_1_t2v_14b_tf_chunkwise_full.yaml b/lightx2v_train/configs/train/tf/wan2_1_t2v_14b_tf_chunkwise_full.yaml index 83cb20977..aa7437b6f 100644 --- a/lightx2v_train/configs/train/tf/wan2_1_t2v_14b_tf_chunkwise_full.yaml +++ b/lightx2v_train/configs/train/tf/wan2_1_t2v_14b_tf_chunkwise_full.yaml @@ -2,6 +2,7 @@ model: name: wan_t2v_14b pretrained_model_name_or_path: /path/to/models/Wan2.1-T2V-14B/ causal: true + num_frame_per_chunk: 3 max_sequence_length: 512 running_dtype: bf16 transformer_param_dtype: fp32 @@ -30,10 +31,16 @@ distributed: data: train: - name: latent_dataset - data_path: /path/to/datasets/causal_forcing_data/dataset/clean_data/ + name: video_dataset + data_path: + - /path/to/datasets/wan_t2v/train.jsonl + video_column: video + prompt_column: prompt + height: 480 + width: 832 + num_frames: 81 + frame_rate: 24 dataset_repeat: 1 - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -43,7 +50,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -72,9 +78,8 @@ training: save_every_iters: 100 save_total_limit: 10 teacher_forcing: - enabled: true mode: chunkwise - num_frame_per_chunk: 3 + num_frame_per_chunk: ${model.num_frame_per_chunk} noise_augmentation_max_timestep: 0 optimizer: learning_rate: 0.000002 diff --git a/lightx2v_train/configs/train/tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml b/lightx2v_train/configs/train/tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml index d257c29bd..8c13273f8 100644 --- a/lightx2v_train/configs/train/tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml +++ b/lightx2v_train/configs/train/tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml @@ -2,6 +2,7 @@ model: name: wan_t2v pretrained_model_name_or_path: /path/to/models/Wan2.1-T2V-1.3B/ causal: true + num_frame_per_chunk: 3 max_sequence_length: 512 running_dtype: bf16 transformer_param_dtype: fp32 @@ -30,10 +31,16 @@ distributed: data: train: - name: latent_dataset - data_path: /path/to/datasets/causal_forcing_data/dataset/clean_data + name: video_dataset + data_path: + - /path/to/datasets/wan_t2v/train.jsonl + video_column: video + prompt_column: prompt + height: 480 + width: 832 + num_frames: 81 + frame_rate: 24 dataset_repeat: 1 - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -43,7 +50,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -72,9 +78,8 @@ training: save_every_iters: 100 save_total_limit: 10 teacher_forcing: - enabled: true mode: chunkwise - num_frame_per_chunk: 3 + num_frame_per_chunk: ${model.num_frame_per_chunk} noise_augmentation_max_timestep: 0 optimizer: learning_rate: 0.000002 diff --git a/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_cache_96f_sp4.yaml b/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_cache_96f_sp4.yaml deleted file mode 100644 index 78a4a8abc..000000000 --- a/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_cache_96f_sp4.yaml +++ /dev/null @@ -1,118 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - causal: true - latent_channels: 48 - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: fp32 - vae_dtype: fp32 - t5_dtype: bf16 - t5_cpu: true - load_vae: false - load_text_encoder: false - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: true - size: 4 - fsdp2: - enabled: true - size: 1 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: /path/to/datasets/cache_data/wan2_2_ti2v_5b_tf_96f - dataset_repeat: 1 - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: teacher_forcing - train_type: full - max_train_iters: 5000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 2 - teacher_forcing: - enabled: true - mode: chunkwise - num_frame_per_chunk: 8 - noise_augmentation_max_timestep: 0 - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: ./output_train/wan2_2_ti2v_5b_tf_cache_96f_sp4 - -inference: - method: wan_ti2v_5b_ar_infer - default_width: 1280 - default_height: 704 - num_frames: 61 - num_inference_steps: 50 - fps: 24 - enable_cfg: true - cfg_guidance_scale: 3.0 - seed: 42 - save_latents_only: true - output_dir: ./output_infer/wan2_2_ti2v_5b_tf_cache_96f_sp4 - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_chunkwise_full_sp4.yaml b/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_chunkwise_full_sp4.yaml deleted file mode 100644 index dc88c14f1..000000000 --- a/lightx2v_train/configs/train/tf/wan2_2_ti2v_5b_tf_chunkwise_full_sp4.yaml +++ /dev/null @@ -1,116 +0,0 @@ -model: - name: wan_ti2v_5b - pretrained_model_name_or_path: /path/to/models/wan2.2/Wan2.2-TI2V-5B - causal: true - latent_channels: 48 - max_sequence_length: 512 - running_dtype: bf16 - transformer_param_dtype: bf16 - vae_dtype: fp32 - t5_dtype: bf16 - load_vae: false - load_text_encoder: false - load_transformer: true - vae_stride: - - 4 - - 16 - - 16 - patch_size: - - 1 - - 2 - - 2 - -distributed: - backend: nccl - timeout_minutes: 60 - sequence_parallel: - enabled: true - size: 4 - fsdp2: - enabled: true - size: 1 - reshard_after_forward: - root_reshard: false - block_reshard: true - mixed_precision: - param_dtype: bf16 - reduce_dtype: fp32 - output_dtype: null - cast_forward_inputs: false - -data: - train: - name: latent_dataset - data_path: ./cache_data/wan2_2_ti2v_5b_tf_96f - dataset_repeat: 1 - batch_size: 1 - num_workers: 8 - pin_memory: true - shuffle: true - drop_last: true - val: - name: prompt_dataset - data_path: - - /path/to/prompts/wan_t2v_prompts.txt - max_samples: 1 - batch_size: 1 - num_workers: 0 - shuffle: false - -scheduler: - num_train_timesteps: 1000 - timestep_distribution: uniform - min_t: 0.001 - max_t: 1.0 - time_shift_settings: - do_time_shift: true - shift_type: linear - time_shift_mu: 5.0 - time_shift_power: 1.0 - sigma_min: 0.0 - extra_one_step: true - -training: - method: teacher_forcing - train_type: full - max_train_iters: 5000 - gradient_accumulation_iters: 1 - gradient_checkpointing: true - max_grad_norm: 10.0 - lr_scheduler: constant - lr_warmup_iters: 0 - save_every_iters: 100 - save_total_limit: 2 - teacher_forcing: - enabled: true - mode: chunkwise - num_frame_per_chunk: 8 - noise_augmentation_max_timestep: 0 - optimizer: - learning_rate: 0.000002 - adam_beta1: 0.0 - adam_beta2: 0.999 - weight_decay: 0.01 - adam_epsilon: 0.00000001 - output_dir: ./output_train/wan2_2_ti2v_5b_tf_chunkwise_full_sp4 - -inference: - method: wan_ti2v_5b_infer - default_width: 1280 - default_height: 704 - num_frames: 381 - num_inference_steps: 50 - fps: 24 - enable_cfg: true - cfg_guidance_scale: 3.0 - seed: 42 - output_dir: ./output_infer/wan2_2_ti2v_5b_tf_chunkwise_full_sp4 - infer_every_iters: null - -logging: - rank_zero_only: true - train_log_every_iters: 10 - infer_log_every_steps: 10 - -resume: - auto_resume: true diff --git a/lightx2v_train/data_process/README.md b/lightx2v_train/data_process/README.md deleted file mode 100644 index 5464f1e01..000000000 --- a/lightx2v_train/data_process/README.md +++ /dev/null @@ -1,296 +0,0 @@ -# 训练数据与预处理 - -本文说明训练数据格式，以及 WAN 和 LTX2 的 latent 数据集构建方法。预处理结果包含统一的 `metadata.jsonl`，可由 `latent_dataset` 直接读取。 - -```text -data_process/ -├── wan/build_latent_dataset.py -└── ltx/build_latent_dataset.py -``` - -## 数据集选择 - -| 训练方式 | 数据集 | 输入 | -| --- | --- | --- | -| Flow | `video_dataset` | 视频路径和 caption 元数据，训练时在线编码 | -| Flow | `latent_dataset` | 预先生成的视频、音频和文本缓存 | -| DMD / AR-DMD | `prompt_dataset` | TXT prompt，训练时在线编码 | -| DMD / AR-DMD | `latent_dataset` | 预先生成的文本条件缓存 | -| Teacher Forcing | `latent_dataset` | PT 缓存元数据或 LMDB | - -WAN 和 LTX2 使用相同的数据集入口。所有数据集均返回统一结构： - -```python -{ - "inputs": {}, - "conditioning": {}, - "meta": {}, -} -``` - -- `inputs`：图片、参考图片、视频、视频 latent 和音频 latent 等模型输入。 -- `conditioning`：prompt、正向文本条件和负向文本条件。 -- `meta`：文件路径、分辨率和帧数等样本信息。 - -## video_dataset - -`video_dataset` 读取 JSON、JSONL 或 CSV 元数据，并在训练时加载视频。目录输入必须包含 `metadata.jsonl`。 - -```json -{"video":"videos/000001.mp4","caption":"A person walks through a park."} -``` - -```yaml -data: - train: - name: video_dataset - data_path: /path/to/datasets/metadata.jsonl - video_root: /path/to/datasets/videos - video_column: video - prompt_column: caption - height: 480 - width: 832 - num_frames: 81 - frame_rate: 24 - batch_size: 1 -``` - -视频路径可以是绝对路径，也可以是相对元数据文件或 `video_root` 的路径。 - -## prompt_dataset - -`prompt_dataset` 读取 TXT 或 LIST 文件，每个非空行表示一个 prompt。 - -```yaml -data: - train: - name: prompt_dataset - data_path: /path/to/prompts/train.txt - batch_size: 1 -``` - -## latent_dataset - -`latent_dataset` 读取 JSON、JSONL、CSV 缓存元数据、包含 `metadata.jsonl` 的目录或 LMDB。缓存元数据支持以下字段： - -| 字段 | 含义 | -| --- | --- | -| `video_latent_path` | 视频 latent 的 PT 文件 | -| `audio_latent_path` | 音频 latent 的 PT 文件，LTX2 T2AV 使用 | -| `condition_path` | 正向文本条件的 PT 文件 | -| `negative_condition_path` | 当前样本的负向文本条件，可选 | -| `caption` | 原始文本，可选 | - -WAN 示例： - -```json -{"video_latent_path":"latents/000001.pt","condition_path":"conditions/000001.pt"} -``` - -LTX2 T2AV 示例： - -```json -{"video_latent_path":"latents/000001.pt","audio_latent_path":"audio_latents/000001.pt","condition_path":"conditions/000001.pt"} -``` - -```yaml -data: - train: - name: latent_dataset - data_path: /path/to/datasets/latent_cache - batch_size: 1 -``` - -相对路径以元数据文件所在目录为基准。缓存目录中的 `negative_condition.pt` 会被自动加载；`negative_condition_path` 仅用于指定其他位置的缓存。包含 `data.mdb` 或 `lock.mdb` 的目录会被识别为 LMDB，并支持 WAN causal LMDB 和统一的 `sample_pt` LMDB。 - -常用 DataLoader 配置包括 `dataset_repeat`、`max_samples`、`batch_size`、`num_workers`、`shuffle`、`drop_last` 和 `pin_memory`。 - -## 预处理输入元数据 - -脚本接受 JSON、JSONL 或 CSV。至少需要视频路径和 caption 两列： - -```json -{"video":"videos/000001.mp4","caption":"A person walks through a park."} -``` - -CSV 示例： - -```csv -video,caption -videos/000001.mp4,A person walks through a park. -``` - -通过 `--video-column` 和 `--prompt-column` 指定自定义列名。相对路径可通过 `--video-root` 解析。 - -## WAN - -以下命令同时生成视频 latent、正向文本条件和负向文本条件： - -```bash -cd /path/to/LightX2V - -python lightx2v_train/data_process/wan/build_latent_dataset.py \ - /path/to/datasets/metadata.csv \ - --video-root /path/to/datasets/videos \ - --output-dir /path/to/datasets/wan_latent_cache \ - --model-dir /path/to/models/Wan2.2-TI2V-5B \ - --cache-components all \ - --video-column video \ - --prompt-column caption \ - --height 704 \ - --width 1280 \ - --latent-frames 96 \ - --max-samples 1000 -``` - -`--cache-components` 可选值： - -| 值 | 输出 | -| --- | --- | -| `all` | 视频 latent 和文本条件 | -| `video` | 仅视频 latent | -| `prompt` | 仅文本条件和负向文本条件 | - -输出结构： - -```text -wan_latent_cache/ -├── metadata.jsonl -├── negative_condition.pt -├── latents/ -└── conditions/ -``` - -构建文本缓存时，脚本使用固定的 `WAN_NEGATIVE_PROMPT` 生成 `negative_condition.pt`。仅生成视频 latent 时不会处理文本条件。 - -脚本支持连续片段拼接。文件名形如 `sample_0.mp4`、`sample_1.mp4` 时，会按序拼接到目标帧数；拼接后仍不足目标帧数的样本会被跳过。对应 caption 使用 `--prompt-separator` 连接。 - -## LTX2 - -LTX2 预处理调用官方 LTX-2 工具，生成视频 latent、音频 latent 和文本条件： - -```bash -cd /path/to/LightX2V - -python lightx2v_train/data_process/ltx/build_latent_dataset.py \ - /path/to/datasets/ltx_t2av/metadata.jsonl \ - --output-dir /path/to/datasets/ltx_t2av/cache \ - --model-path /path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors \ - --text-encoder-path /path/to/models/LTX-2 \ - --ltx2-repo /path/to/LTX-2 \ - --resolution-buckets 768x512x241 \ - --batch-size 1 \ - --device cuda -``` - -输出结构： - -```text -cache/ -├── metadata.jsonl -├── negative_condition.pt -├── latents/ -├── audio_latents/ -└── conditions/ -``` - -脚本会先使用固定的 `LTX2_NEGATIVE_PROMPT` 生成 `negative_condition.pt`。该缓存用于 DMD 和 AR-DMD；仅训练 Flow 或 Teacher Forcing 时可通过 `--skip-negative-cache` 跳过。T2AV 的 Flow 和 Teacher Forcing 需要音频 latent；仅处理无音频任务时使用 `--skip-audio`。 - -## 训练 - -训练任务通过配置中的 `data.train` 选择数据集。WAN 和 LTX2 的写法一致。 - -### Flow - -使用原始视频和 caption 时选择 `video_dataset`。视频和文本编码在训练时执行： - -```yaml -data: - train: - name: video_dataset - data_path: /path/to/datasets/metadata.jsonl - video_root: /path/to/datasets/videos - video_column: video - prompt_column: caption - height: 480 - width: 832 - num_frames: 81 - frame_rate: 24 - batch_size: 1 - num_workers: 4 - shuffle: true -``` - -使用预计算缓存时选择 `latent_dataset`： - -```yaml -data: - train: - name: latent_dataset - data_path: /path/to/datasets/latent_cache/metadata.jsonl - batch_size: 1 - num_workers: 4 - shuffle: true -``` - -WAN 的缓存记录包含 `video_latent_path` 和 `condition_path`。LTX2 T2AV 还需要 `audio_latent_path`。 - -### DMD / AR-DMD - -使用 TXT prompt 并在线运行 text encoder 时选择 `prompt_dataset`： - -```yaml -data: - train: - name: prompt_dataset - data_path: /path/to/prompts/train.txt - batch_size: 1 - num_workers: 4 - shuffle: true - drop_last: true -``` - -使用预计算文本条件时选择 `latent_dataset`。元数据中的每条记录需要包含 `condition_path`： - -```yaml -data: - train: - name: latent_dataset - data_path: /path/to/datasets/prompt_cache/metadata.jsonl - batch_size: 1 - num_workers: 4 - shuffle: true - drop_last: true -``` - -`latent_dataset` 会自动加载缓存目录中的 `negative_condition.pt`。 - -### Teacher Forcing - -Teacher Forcing 使用 `latent_dataset`。PT 缓存通过 `metadata.jsonl` 加载： - -```yaml -data: - train: - name: latent_dataset - data_path: /path/to/datasets/tf_cache/metadata.jsonl - batch_size: 1 - num_workers: 4 - shuffle: true - drop_last: true -``` - -LMDB 直接将 `data_path` 指向包含 `data.mdb` 或 `lock.mdb` 的目录： - -```yaml -data: - train: - name: latent_dataset - data_path: /path/to/datasets/tf_cache.lmdb - batch_size: 1 - num_workers: 4 - shuffle: true - drop_last: true -``` - -WAN Teacher Forcing 需要视频 latent 和文本条件。LTX2 T2AV Teacher Forcing 需要视频 latent、音频 latent 和文本条件。 diff --git a/lightx2v_train/data_process/ltx/build_latent_dataset.py b/lightx2v_train/data_process/ltx/build_latent_dataset.py deleted file mode 100644 index 97de832ae..000000000 --- a/lightx2v_train/data_process/ltx/build_latent_dataset.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python3 -"""Build an LTX2 T2AV latent_dataset cache via the LTX preprocessing tools.""" - -import argparse -import json -import os -import sys -from pathlib import Path - -REPO_ROOT = Path(__file__).resolve().parents[2] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from lightx2v_train.data.utils import read_records # noqa: E402 -from lightx2v_train.model_zoo.native.ltx2 import ensure_ltx2_native_path # noqa: E402 -from lightx2v_train.utils.constants import LTX2_NEGATIVE_PROMPT # noqa: E402 - - -def parse_args(): - parser = argparse.ArgumentParser(description="Build LTX2 T2AV video/audio/prompt latent_dataset cache from json/jsonl/csv metadata.") - parser.add_argument("metadata", help="CSV/JSON/JSONL metadata with video and caption columns.") - parser.add_argument("--output-dir", required=True, help="Output cache directory. Writes latents/, audio_latents/, conditions/ and metadata.jsonl.") - parser.add_argument("--model-path", default="/path/to/models/ltx-2.3/ltx-2.3-22b-dev.safetensors") - parser.add_argument("--text-encoder-path", default="/path/to/models/LTX-2") - parser.add_argument( - "--ltx2-repo", - "--ltx2-trainer-repo", - dest="ltx2_repo", - default="/path/to/LTX-2", - help="Reference LTX-2 repo path for process_dataset/process_videos.", - ) - parser.add_argument("--resolution-buckets", default="768x768x49", help="Semicolon separated WxHxF buckets.") - parser.add_argument("--video-column", default=None) - parser.add_argument("--caption-column", "--prompt-column", dest="caption_column", default=None) - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--device", default="cuda") - parser.add_argument("--vae-tiling", action="store_true") - parser.add_argument("--decode", action="store_true") - parser.add_argument("--remove-llm-prefixes", action="store_true") - parser.add_argument("--skip-audio", action="store_true", help="Forwarded to LTX preprocessing. T2AV flow/tf usually need audio.") - parser.add_argument("--audio-durations", default=None, help="Audio duration buckets, e.g. 2.0;4.0;8.0.") - parser.add_argument("--load-text-encoder-in-8bit", action="store_true") - parser.add_argument("--skip-negative-cache", action="store_true") - parser.add_argument("--max-samples", type=int, default=None) - parser.add_argument("--overwrite", action="store_true") - return parser.parse_args() - - -def add_ltx2_paths(ltx2_repo): - ensure_ltx2_native_path() - repo = Path(ltx2_repo).expanduser().resolve() - paths = [ - repo / "packages" / "ltx-trainer" / "src", - repo / "packages" / "ltx-trainer" / "scripts", - ] - missing = [str(path) for path in paths if not path.exists()] - if missing: - raise FileNotFoundError(f"Missing LTX-2 trainer paths: {missing}") - for path in reversed(paths): - path_str = str(path) - if path_str not in sys.path: - sys.path.insert(0, path_str) - - -def atomic_write_jsonl(path, rows): - path = Path(path) - tmp_path = path.with_suffix(f"{path.suffix}.tmp.{os.getpid()}") - with tmp_path.open("w", encoding="utf-8") as handle: - for row in rows: - handle.write(json.dumps(row, ensure_ascii=False) + "\n") - os.replace(tmp_path, path) - - -def write_manifest(root, require_audio): - root = Path(root).expanduser().resolve() - video_dir = root / "latents" - audio_dir = root / "audio_latents" - condition_dir = root / "conditions" - required_dirs = [video_dir, condition_dir] - if require_audio: - required_dirs.append(audio_dir) - for path in required_dirs: - if not path.is_dir(): - raise FileNotFoundError(f"Missing latent dataset directory: {path}") - - audio_files = {path.relative_to(audio_dir): path for path in audio_dir.rglob("*.pt")} if audio_dir.is_dir() else {} - condition_files = {path.relative_to(condition_dir): path for path in condition_dir.rglob("*.pt")} - rows = [] - skipped = 0 - for video_path in sorted(video_dir.rglob("*.pt")): - rel_path = video_path.relative_to(video_dir) - condition_rel_path = rel_path - if rel_path.name.startswith("latent_"): - condition_rel_path = rel_path.parent / f"condition_{rel_path.name[7:]}" - audio_path = audio_files.get(rel_path) - condition_path = condition_files.get(condition_rel_path) - if condition_path is None or (require_audio and audio_path is None): - skipped += 1 - continue - row = { - "video_latent_path": str(video_path.relative_to(root)), - "condition_path": str(condition_path.relative_to(root)), - } - if audio_path is not None: - row["audio_latent_path"] = str(audio_path.relative_to(root)) - rows.append(row) - - manifest_path = root / "metadata.jsonl" - atomic_write_jsonl(manifest_path, rows) - return manifest_path, len(rows), skipped - - -def write_negative_condition(args, output_dir, compute_captions_embeddings): - negative_path = output_dir / "negative_condition.pt" - if negative_path.is_file() and not args.overwrite: - return negative_path - - metadata_path = output_dir / f"_negative_prompt.{os.getpid()}.jsonl" - atomic_write_jsonl(metadata_path, [{"video": negative_path.name, "caption": LTX2_NEGATIVE_PROMPT}]) - try: - compute_captions_embeddings( - dataset_file=str(metadata_path), - output_dir=str(output_dir), - model_path=args.model_path, - text_encoder_path=args.text_encoder_path, - caption_column="caption", - media_column="video", - batch_size=1, - device=args.device, - load_in_8bit=args.load_text_encoder_in_8bit, - overwrite=args.overwrite, - ) - finally: - metadata_path.unlink(missing_ok=True) - if not negative_path.is_file(): - raise RuntimeError(f"Failed to build negative prompt cache: {negative_path}") - return negative_path - - -def limited_metadata(metadata, output_dir, max_samples, caption_column): - if max_samples is None: - return metadata - rows = [] - for row in read_records(metadata, prompt_column=caption_column or "caption"): - rows.append(row) - if len(rows) >= max_samples: - break - if not rows: - raise RuntimeError(f"No rows found in metadata: {metadata}") - limited_path = Path(output_dir) / "_limited_metadata.jsonl" - atomic_write_jsonl(limited_path, rows) - return str(limited_path) - - -def main(): - args = parse_args() - add_ltx2_paths(args.ltx2_repo) - - from ltx_trainer.gpu_utils import free_gpu_memory_context - from process_captions import compute_captions_embeddings - from process_dataset import preprocess_dataset - from process_videos import parse_resolution_buckets - - output_dir = Path(args.output_dir).expanduser().resolve() - output_dir.mkdir(parents=True, exist_ok=True) - if not args.skip_negative_cache: - with free_gpu_memory_context(): - negative_path = write_negative_condition(args, output_dir, compute_captions_embeddings) - print(f"Wrote negative condition: {negative_path}", flush=True) - - metadata = limited_metadata(args.metadata, output_dir, args.max_samples, args.caption_column) - parsed_buckets = parse_resolution_buckets(args.resolution_buckets) if args.resolution_buckets else None - audio_durations = [float(item) for item in args.audio_durations.split(";")] if args.audio_durations else None - - preprocess_dataset( - dataset_file=metadata, - resolution_buckets=parsed_buckets, - model_path=args.model_path, - text_encoder_path=args.text_encoder_path, - device=args.device, - output_dir=str(output_dir), - video_column=args.video_column, - caption_column=args.caption_column, - batch_size=args.batch_size, - vae_tiling=args.vae_tiling, - decode=args.decode, - remove_llm_prefixes=args.remove_llm_prefixes, - skip_audio=args.skip_audio, - audio_durations=audio_durations, - load_text_encoder_in_8bit=args.load_text_encoder_in_8bit, - overwrite=args.overwrite, - ) - manifest_path, count, skipped = write_manifest(output_dir, require_audio=not args.skip_audio) - print(f"Wrote latent manifest: {manifest_path} ({count} samples, {skipped} skipped)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/lightx2v_train/data_process/minimax_h3/build_minimax_h3_prompt_latents.py b/lightx2v_train/data_process/minimax_h3/build_minimax_h3_prompt_latents.py deleted file mode 100644 index e19a557c0..000000000 --- a/lightx2v_train/data_process/minimax_h3/build_minimax_h3_prompt_latents.py +++ /dev/null @@ -1,134 +0,0 @@ -#!/usr/bin/env python3 -"""Encode one-prompt-per-line text into an H3 latent_dataset cache.""" - -import argparse -import json -import os -from pathlib import Path - -import torch - -TEXT_ENCODER_LAYER = 50 -TEXT_TAG = 1 - - -def parse_args(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("prompts", help="UTF-8 txt file; each non-empty line is one prompt.") - parser.add_argument("--model-path", default="/data/nvme6/gushiqiao/models/MiniMax-H3") - parser.add_argument("--output-dir", required=True) - parser.add_argument("--device", default="cuda:0") - parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16") - parser.add_argument("--max-samples", type=int) - parser.add_argument("--start-index", type=int, default=0) - parser.add_argument("--overwrite", action="store_true") - return parser.parse_args() - - -def read_prompts(path, start_index=0, max_samples=None): - with Path(path).expanduser().open("r", encoding="utf-8") as handle: - prompts = [line.strip() for line in handle if line.strip()] - prompts = prompts[start_index:] - if max_samples is not None: - prompts = prompts[:max_samples] - if not prompts: - raise RuntimeError(f"No non-empty prompts found in {path}.") - return prompts - - -def atomic_write_jsonl(path, rows): - path = Path(path) - temporary = path.with_suffix(f"{path.suffix}.tmp.{os.getpid()}") - with temporary.open("w", encoding="utf-8") as handle: - for row in rows: - handle.write(json.dumps(row, ensure_ascii=False) + "\n") - os.replace(temporary, path) - - -def load_conditioner(model_path, device, dtype): - try: - from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor - except ImportError as exc: - raise ImportError("MiniMax-H3 prompt encoding requires a Transformers build with Qwen3-VL support. The model's local_diffusers environment contains the required version.") from exc - - root = Path(model_path).expanduser().resolve() - text_encoder_path = root / "text_encoder" - tokenizer_path = root / "tokenizer" - processor_path = root / "processor" - for path in (text_encoder_path, tokenizer_path, processor_path): - if not path.exists(): - raise FileNotFoundError(f"Missing MiniMax-H3 component: {path}") - - load_kwargs = {"dtype": dtype, "local_files_only": True} - if str(device) != "cpu": - load_kwargs["device_map"] = {"": str(device)} - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(str(text_encoder_path), **load_kwargs) - text_encoder.requires_grad_(False).eval() - tokenizer = Qwen2TokenizerFast.from_pretrained(str(tokenizer_path), local_files_only=True) - processor = Qwen3VLProcessor.from_pretrained(str(processor_path), local_files_only=True) - return text_encoder, tokenizer, processor - - -@torch.inference_mode() -def encode_prompt(text_encoder, tokenizer, processor, prompt, device, output_dtype): - token_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] - if not token_ids: - raise ValueError("The tokenizer produced no tokens for a non-empty prompt.") - input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) - mm_token_type_ids = torch.tensor(processor.create_mm_token_type_ids([token_ids]), dtype=torch.long, device=device) - outputs = text_encoder.model( - input_ids=input_ids, - attention_mask=torch.ones_like(input_ids), - mm_token_type_ids=mm_token_type_ids, - pixel_values=None, - image_grid_thw=None, - use_cache=False, - output_hidden_states=True, - ) - if len(outputs.hidden_states) <= TEXT_ENCODER_LAYER: - raise RuntimeError(f"Qwen3-VL returned {len(outputs.hidden_states)} hidden states; H3 requires index {TEXT_ENCODER_LAYER}.") - prompt_embeds = outputs.hidden_states[TEXT_ENCODER_LAYER].to(dtype=output_dtype).cpu() - text_token_tags = torch.full((len(token_ids),), TEXT_TAG, dtype=torch.long) - return {"prompt_embeds": prompt_embeds, "text_token_tags": text_token_tags} - - -def main(): - args = parse_args() - dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.dtype] - prompts = read_prompts(args.prompts, args.start_index, args.max_samples) - output_dir = Path(args.output_dir).expanduser().resolve() - condition_dir = output_dir / "conditions" - condition_dir.mkdir(parents=True, exist_ok=True) - text_encoder, tokenizer, processor = load_conditioner(args.model_path, args.device, dtype) - - rows = [] - for local_index, prompt in enumerate(prompts): - source_index = args.start_index + local_index - relative_path = Path("conditions") / f"condition_{source_index:08d}.pt" - output_path = output_dir / relative_path - if args.overwrite or not output_path.is_file(): - condition = encode_prompt(text_encoder, tokenizer, processor, prompt, args.device, dtype) - torch.save( - { - "conditioning": {"positive": condition}, - "prompt": prompt, - "source_index": source_index, - }, - output_path, - ) - rows.append( - { - "id": source_index, - "caption": prompt, - "condition_path": str(relative_path), - } - ) - print(f"[{local_index + 1}/{len(prompts)}] {output_path}", flush=True) - - manifest = output_dir / "metadata.jsonl" - atomic_write_jsonl(manifest, rows) - print(f"Wrote {len(rows)} samples to {manifest}", flush=True) - - -if __name__ == "__main__": - main() diff --git a/lightx2v_train/data_process/wan/build_latent_dataset.py b/lightx2v_train/data_process/wan/build_latent_dataset.py deleted file mode 100644 index 73b22d663..000000000 --- a/lightx2v_train/data_process/wan/build_latent_dataset.py +++ /dev/null @@ -1,549 +0,0 @@ -#!/usr/bin/env python3 -"""Build a Wan latent_dataset cache from video/caption metadata. - -The output is directly consumable by lightx2v_train.data.latent_dataset: - - output_dir/ - metadata.jsonl - latents/*.pt - conditions/*.pt - -Each manifest row may contain video_latent_path, condition_path, or both, -depending on --cache-components. -""" - -from __future__ import annotations - -import argparse -import json -import math -import os -import re -import sys -from dataclasses import dataclass -from pathlib import Path - -import imageio.v2 as imageio -import numpy as np -import torch -from PIL import Image - -REPO_ROOT = Path(__file__).resolve().parents[2] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from lightx2v_train.data.utils import prompt_text, read_records, record_value, resolve_data_path # noqa: E402 -from lightx2v_train.model_zoo.native.wan.modules.t5 import T5EncoderModel # noqa: E402 -from lightx2v_train.model_zoo.native.wan.modules.vae2_2 import Wan2_2_VAE # noqa: E402 -from lightx2v_train.utils.constants import WAN_NEGATIVE_PROMPT # noqa: E402 - -TORCH_DTYPES = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -@dataclass -class ClipEntry: - path: Path - prompt: str - frame_count: int | None - metadata_path: Path - row_index: int - prefix: str - number: int | None - - -def parse_args(): - parser = argparse.ArgumentParser(description="Build Wan video/prompt latent_dataset cache from json/jsonl/csv metadata.") - parser.add_argument("metadata", nargs="+", help="CSV/JSON/JSONL metadata file(s), or video directory.") - parser.add_argument("--output-dir", required=True, help="Output latent_dataset directory.") - parser.add_argument("--cache-components", default="all", choices=["all", "video", "prompt"], help="Build video latents, prompt conditions, or both.") - parser.add_argument("--model-dir", default="/path/to/models/Wan2.2-TI2V-5B", help="Wan2.2 model directory.") - parser.add_argument("--vae-checkpoint", default=None, help="Override Wan VAE checkpoint path.") - parser.add_argument("--t5-checkpoint", default=None, help="Override T5 checkpoint path.") - parser.add_argument("--tokenizer-path", default=None, help="Override tokenizer path.") - parser.add_argument("--video-root", action="append", default=[], help="Root used to resolve relative video paths. Can be passed more than once.") - parser.add_argument("--video-column", default="video") - parser.add_argument("--prompt-column", default="caption") - parser.add_argument("--height", type=int, default=704) - parser.add_argument("--width", type=int, default=1280) - parser.add_argument("--raw-frames", type=int, default=None, help="Raw frames per sample. Defaults to 1 + ratio * (latent_frames - 1).") - parser.add_argument("--latent-frames", type=int, default=96) - parser.add_argument("--temporal-compression-ratio", type=int, default=4) - parser.add_argument("--frame-rate", type=float, default=24.0) - parser.add_argument("--fix-frame-rate", action="store_true") - parser.add_argument("--prompt-separator", default=" The scene transitions. ") - parser.add_argument("--video-glob", default="*.mp4", help="Glob used when metadata path is a directory.") - parser.add_argument("--recursive", action="store_true", help="Recursively scan video directories.") - parser.add_argument("--max-samples", type=int, default=None) - parser.add_argument("--max-groups", type=int, default=None) - parser.add_argument("--rank", type=int, default=int(os.environ.get("RANK", "0"))) - parser.add_argument("--local-rank", "--local_rank", dest="local_rank", type=int, default=int(os.environ.get("LOCAL_RANK", "0"))) - parser.add_argument("--world-size", type=int, default=int(os.environ.get("WORLD_SIZE", "1"))) - parser.add_argument("--device", default="cuda") - parser.add_argument("--text-device", default=None, help="Defaults to --device.") - parser.add_argument("--vae-dtype", default="fp32", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--t5-dtype", default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--save-dtype", default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--max-sequence-length", type=int, default=512) - parser.add_argument("--vae-chunk-latent-frames", type=int, default=16, help="Use 0 to encode the full video in one VAE call.") - parser.add_argument("--vae-halo-latents", type=int, default=28) - parser.add_argument("--skip-negative-cache", action="store_true") - parser.add_argument("--skip-missing", action="store_true", default=True) - parser.add_argument("--no-skip-missing", dest="skip_missing", action="store_false") - parser.add_argument("--strict-decode", action="store_true") - parser.add_argument("--overwrite", action="store_true") - parser.add_argument("--dry-run", action="store_true") - parser.add_argument("--log-every", type=int, default=1) - parser.add_argument("--log-source-limit", type=int, default=0) - return parser.parse_args() - - -def need_video(args): - return args.cache_components in {"all", "video"} - - -def need_prompt(args): - return args.cache_components in {"all", "prompt"} - - -def crop_and_resize(image, target_height, target_width): - width, height = image.size - scale = max(target_width / width, target_height / height) - resized_width = round(width * scale) - resized_height = round(height * scale) - image = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR) - left = max(0, (resized_width - target_width) // 2) - top = max(0, (resized_height - target_height) // 2) - return image.crop((left, top, left + target_width, top + target_height)) - - -def frame_to_tensor(frame): - array = np.array(frame, dtype=np.uint8, copy=True) - return torch.from_numpy(array).permute(2, 0, 1).contiguous() - - -def numeric_parts(path): - match = re.search(r"(\d+)$", path.stem) - if not match: - return path.stem, None - return path.stem[: match.start(1)], int(match.group(1)) - - -def load_sidecar_prompt(video_path, prompt_column): - for sidecar in (video_path.with_suffix(".txt"), video_path.with_suffix(".json")): - if not sidecar.is_file(): - continue - try: - if sidecar.suffix == ".txt": - return sidecar.read_text(encoding="utf-8").strip() - payload = json.loads(sidecar.read_text(encoding="utf-8")) - if isinstance(payload, dict): - return str(payload.get(prompt_column, payload.get("caption", payload.get("prompt", "")))) - except Exception: - return "" - return "" - - -def iter_rows(metadata_path, args): - metadata_path = Path(metadata_path) - if metadata_path.is_dir(): - pattern = f"**/{args.video_glob}" if args.recursive else args.video_glob - for index, video_path in enumerate(sorted(metadata_path.glob(pattern))): - yield {"video": str(video_path), "caption": load_sidecar_prompt(video_path, args.prompt_column)}, index - return - for index, row in enumerate(read_records(metadata_path, prompt_column=args.prompt_column)): - yield row, index - - -def resolve_video(metadata_path, row, args): - video_value = record_value(row, args.video_column, "video_path", "video") - if video_value is None: - return None - base_dir = metadata_path if metadata_path.is_dir() else metadata_path.parent - return resolve_data_path(video_value, base_dir, roots=args.video_root, subdirs=("video", "videos")) - - -def parse_optional_int(value): - if value is None or value == "": - return None - try: - return int(float(value)) - except (TypeError, ValueError): - return None - - -def load_entries(args): - entries = [] - entry_limit = None - if not need_video(args) and args.max_samples is not None: - entry_limit = int(args.max_samples) * max(1, int(args.world_size)) - for metadata in args.metadata: - metadata_path = Path(metadata) - for row, row_index in iter_rows(metadata_path, args): - video_path = resolve_video(metadata_path, row, args) - if need_video(args) and args.skip_missing and (video_path is None or not video_path.is_file()): - continue - prompt = prompt_text(row, args.prompt_column) - if not need_video(args): - entries.append(ClipEntry(path=video_path or Path(""), prompt=prompt, frame_count=None, metadata_path=metadata_path, row_index=row_index, prefix="", number=None)) - if entry_limit is not None and len(entries) >= entry_limit: - return entries - continue - prefix, number = numeric_parts(video_path) - entries.append( - ClipEntry( - path=video_path, - prompt=prompt, - frame_count=parse_optional_int(record_value(row, "frames", "frame", "num_frames")), - metadata_path=metadata_path, - row_index=row_index, - prefix=prefix, - number=number, - ) - ) - return entries - - -def group_consecutive(entries): - entries = sorted( - entries, - key=lambda item: ( - str(item.path.parent), - item.prefix, - item.number is None, - item.number if item.number is not None else item.path.stem, - item.path.name, - ), - ) - groups = [] - current = [] - for entry in entries: - if not current: - current = [entry] - continue - prev = current[-1] - is_next = prev.path.parent == entry.path.parent and prev.prefix == entry.prefix and prev.number is not None and entry.number is not None and entry.number == prev.number + 1 - if is_next: - current.append(entry) - else: - groups.append(current) - current = [entry] - if current: - groups.append(current) - return groups - - -def usable_frame_count(reader, args): - total_raw_frames = int(reader.count_frames()) - if not args.fix_frame_rate: - return total_raw_frames - metadata = reader.get_meta_data() - duration = metadata.get("duration") or total_raw_frames / metadata.get("fps", args.frame_rate) - return int(math.floor(duration * args.frame_rate)) - - -def iter_frame_ids(reader, args): - total_raw_frames = int(reader.count_frames()) - raw_fps = reader.get_meta_data().get("fps", args.frame_rate) - count = usable_frame_count(reader, args) - for frame_id in range(count): - if args.fix_frame_rate: - target_time = frame_id / args.frame_rate - yield min(int(round(target_time * raw_fps)), total_raw_frames - 1) - else: - yield frame_id - - -def iter_video_frames(video_path, args): - reader = imageio.get_reader(video_path) - try: - for frame_id in iter_frame_ids(reader, args): - frame = Image.fromarray(reader.get_data(frame_id)).convert("RGB") - frame = crop_and_resize(frame, args.height, args.width) - yield frame_to_tensor(frame) - finally: - reader.close() - - -def concat_prompt(parts, separator): - cleaned = [part.strip() for part in parts if part and part.strip()] - return separator.join(cleaned) if cleaned else " " - - -def iter_group_windows(group, args): - frames = [] - prompt_parts = [] - source_paths = [] - source_set = set() - window_index = 0 - - def reset_window(): - frames.clear() - prompt_parts.clear() - source_paths.clear() - source_set.clear() - - for clip in group: - try: - for frame in iter_video_frames(clip.path, args): - path_key = str(clip.path) - if path_key not in source_set: - source_set.add(path_key) - source_paths.append(path_key) - prompt_parts.append(clip.prompt) - frames.append(frame) - if len(frames) == args.raw_frames: - video = torch.stack(frames, dim=1).contiguous() - yield window_index, video, concat_prompt(prompt_parts, args.prompt_separator), list(source_paths) - window_index += 1 - reset_window() - except Exception as error: - if args.strict_decode: - raise - print(f"[warn] skip bad video {clip.path}: {error}", flush=True) - reset_window() - - -def count_group_frames(group, args): - if all(clip.frame_count is not None for clip in group): - return sum(clip.frame_count for clip in group) - total = 0 - for clip in group: - try: - reader = imageio.get_reader(clip.path) - try: - total += usable_frame_count(reader, args) - finally: - reader.close() - except Exception as error: - if args.strict_decode: - raise - print(f"[warn] cannot count {clip.path}: {error}", flush=True) - return total - - -def latent_range_to_raw_window(latent_start, latent_end, ratio): - if latent_start == 0: - return 0, 1 + ratio * (latent_end - 1), 0 - return ratio * (latent_start - 1), 1 + ratio * (latent_end - 1), 1 - - -def encode_video_latent(vae, video_cpu, args): - device = torch.device(args.device) - vae_dtype = TORCH_DTYPES[args.vae_dtype] - latent_parts = [] - - def prepare_video_chunk(chunk_cpu): - return chunk_cpu.to(device=device, dtype=vae_dtype).div_(127.5).sub_(1.0) - - with torch.inference_mode(): - if args.vae_chunk_latent_frames <= 0: - video = prepare_video_chunk(video_cpu) - return vae.encode([video])[0].cpu() - chunk = args.vae_chunk_latent_frames - for keep_start in range(0, args.latent_frames, chunk): - keep_end = min(args.latent_frames, keep_start + chunk) - halo_start = max(0, keep_start - args.vae_halo_latents) - raw_start, raw_end, pseudo_prefix = latent_range_to_raw_window(halo_start, keep_end, args.temporal_compression_ratio) - drop_latents = pseudo_prefix + (keep_start - halo_start) - video = prepare_video_chunk(video_cpu[:, raw_start:raw_end].contiguous()) - encoded = vae.encode([video])[0].cpu() - keep = encoded[:, drop_latents : drop_latents + (keep_end - keep_start)].contiguous() - if keep.shape[1] != keep_end - keep_start: - raise RuntimeError(f"VAE chunk produced {keep.shape[1]} latent frames, expected {keep_end - keep_start}.") - latent_parts.append(keep) - del video, encoded, keep - if device.type == "cuda": - torch.cuda.empty_cache() - return torch.cat(latent_parts, dim=1) - - -def encode_prompt_embed(text_encoder, prompt, args): - device = torch.device(args.text_device or args.device) - with torch.inference_mode(): - contexts = text_encoder([prompt], device) - context = contexts[0][: args.max_sequence_length].to(dtype=TORCH_DTYPES[args.save_dtype]).cpu() - if context.shape[0] < args.max_sequence_length: - pad = context.new_zeros(args.max_sequence_length - context.shape[0], context.shape[1]) - context = torch.cat([context, pad], dim=0) - return context - - -def load_vae(args): - device = torch.device(args.device) - if device.type == "cuda": - torch.cuda.set_device(device) - model_dir = Path(args.model_dir) - checkpoint = args.vae_checkpoint or str(model_dir / "Wan2.2_VAE.pth") - print(f"loading Wan VAE: {checkpoint}", flush=True) - return Wan2_2_VAE(vae_pth=checkpoint, dtype=TORCH_DTYPES[args.vae_dtype], device=device) - - -def load_text_encoder(args): - model_dir = Path(args.model_dir) - checkpoint = args.t5_checkpoint or str(model_dir / "models_t5_umt5-xxl-enc-bf16.pth") - tokenizer_path = args.tokenizer_path or str(model_dir / "google" / "umt5-xxl") - device = torch.device(args.text_device or args.device) - print(f"loading Wan T5: {checkpoint}", flush=True) - return T5EncoderModel( - text_len=args.max_sequence_length, - dtype=TORCH_DTYPES[args.t5_dtype], - device=device, - checkpoint_path=checkpoint, - tokenizer_path=tokenizer_path, - ) - - -def save_pt(path, payload, overwrite): - if path.exists() and not overwrite: - return False - path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") - torch.save(payload, tmp_path) - tmp_path.replace(path) - return True - - -def condition_payload(prompt_embed, prompt): - return {"prompt": prompt, "conditions": {"prompt_embed": prompt_embed.contiguous()}} - - -def validate_args(args): - expected_raw = 1 + args.temporal_compression_ratio * (args.latent_frames - 1) - if args.raw_frames != expected_raw: - raise ValueError(f"--raw-frames must be 1 + ratio * (latent_frames - 1) = {expected_raw}, got {args.raw_frames}.") - if args.height % 16 != 0 or args.width % 16 != 0: - raise ValueError(f"height/width must be divisible by 16, got {args.height}x{args.width}.") - if args.rank < 0 or args.rank >= args.world_size: - raise ValueError(f"rank must be in [0, world_size), got rank={args.rank}, world_size={args.world_size}.") - - -def manifest_path(output_dir, args): - if args.world_size == 1: - return output_dir / "metadata.jsonl" - return output_dir / f"metadata_rank{args.rank:02d}.jsonl" - - -def build_prompt_only(entries, args, text_encoder, output_dir): - rows = [] - conditions_dir = output_dir / "conditions" - save_dtype = TORCH_DTYPES[args.save_dtype] - for idx, entry in enumerate(entries): - if idx % args.world_size != args.rank: - continue - if args.max_samples is not None and len(rows) >= args.max_samples: - break - condition_rel = Path("conditions") / f"condition_{idx:08d}.pt" - condition_path = output_dir / condition_rel - prompt_embed = encode_prompt_embed(text_encoder, entry.prompt, args).to(dtype=save_dtype) - save_pt(condition_path, condition_payload(prompt_embed, entry.prompt), args.overwrite) - rows.append({"condition_path": str(condition_rel), "caption": entry.prompt, "row_index": entry.row_index}) - return rows - - -def main(): - args = parse_args() - if args.raw_frames is None: - args.raw_frames = 1 + args.temporal_compression_ratio * (args.latent_frames - 1) - if args.device == "cuda": - args.device = f"cuda:{args.local_rank}" - if args.text_device == "cuda": - args.text_device = f"cuda:{args.local_rank}" - validate_args(args) - output_dir = Path(args.output_dir).expanduser().resolve() - output_dir.mkdir(parents=True, exist_ok=True) - - entries = load_entries(args) - print(f"entries={len(entries)} cache_components={args.cache_components} output_dir={output_dir}", flush=True) - - if args.dry_run and not need_video(args): - print(f"dry_run_prompt_samples={len(entries)}", flush=True) - return - - if need_video(args): - groups = group_consecutive(entries) - if args.max_groups is not None: - groups = groups[: args.max_groups] - local_groups = [(index, group) for index, group in enumerate(groups) if index % args.world_size == args.rank] - print(f"groups={len(groups)} local_groups={len(local_groups)} raw_frames={args.raw_frames} latent_frames={args.latent_frames}", flush=True) - if args.dry_run: - total_windows = 0 - for group_index, group in local_groups: - frame_count = count_group_frames(group, args) - windows = frame_count // args.raw_frames - total_windows += windows - print(f"group={group_index:06d} clips={len(group)} frames={frame_count} windows={windows} first={group[0].path} last={group[-1].path}", flush=True) - print(f"dry_run_windows={total_windows}", flush=True) - return - else: - groups = [] - local_groups = [] - - vae = load_vae(args) if need_video(args) else None - text_encoder = load_text_encoder(args) if need_prompt(args) else None - save_dtype = TORCH_DTYPES[args.save_dtype] - rows = [] - - if need_prompt(args) and not args.skip_negative_cache: - negative_path = output_dir / "negative_condition.pt" - prompt_embed = encode_prompt_embed(text_encoder, WAN_NEGATIVE_PROMPT, args).to(dtype=save_dtype) - save_pt(negative_path, condition_payload(prompt_embed, WAN_NEGATIVE_PROMPT), args.overwrite) - - if not need_video(args): - rows = build_prompt_only(entries, args, text_encoder, output_dir) - else: - saved = 0 - for local_index, (group_index, group) in enumerate(local_groups): - if args.log_every > 0 and local_index % args.log_every == 0: - print(f"processing group {local_index + 1}/{len(local_groups)} global={group_index}", flush=True) - for window_index, video, prompt, source_paths in iter_group_windows(group, args): - if args.max_samples is not None and saved >= args.max_samples: - break - stem = f"sample_g{group_index:06d}_w{window_index:03d}" - row = { - "caption": prompt, - "video_paths": source_paths, - "raw_num_frames": args.raw_frames, - "latent_num_frames": args.latent_frames, - "height": args.height, - "width": args.width, - } - if need_video(args): - latent_rel = Path("latents") / f"{stem}.pt" - latent = encode_video_latent(vae, video, args) - if latent.shape[1] != args.latent_frames: - raise RuntimeError(f"Expected {args.latent_frames} latent frames, got shape={tuple(latent.shape)}") - latent_payload = { - "latents": latent.to(dtype=save_dtype).contiguous(), - "num_frames": int(latent.shape[1]), - "height": int(latent.shape[2]), - "width": int(latent.shape[3]), - "fps": args.frame_rate, - } - save_pt(output_dir / latent_rel, latent_payload, args.overwrite) - row["video_latent_path"] = str(latent_rel) - if need_prompt(args): - condition_rel = Path("conditions") / f"{stem}.pt" - prompt_embed = encode_prompt_embed(text_encoder, prompt, args).to(dtype=save_dtype) - save_pt(output_dir / condition_rel, condition_payload(prompt_embed, prompt), args.overwrite) - row["condition_path"] = str(condition_rel) - rows.append(row) - saved += 1 - del video - if args.max_samples is not None and saved >= args.max_samples: - break - - out_manifest = manifest_path(output_dir, args) - tmp_manifest = out_manifest.with_suffix(f"{out_manifest.suffix}.tmp.{os.getpid()}") - with tmp_manifest.open("w", encoding="utf-8") as handle: - for row in rows: - handle.write(json.dumps(row, ensure_ascii=False) + "\n") - os.replace(tmp_manifest, out_manifest) - print(f"Wrote manifest: {out_manifest} ({len(rows)} samples)", flush=True) - - -if __name__ == "__main__": - main() diff --git a/lightx2v_train/infer.py b/lightx2v_train/infer.py index d0e1c9ed2..ef8308933 100644 --- a/lightx2v_train/infer.py +++ b/lightx2v_train/infer.py @@ -5,7 +5,7 @@ import torch from loguru import logger -from lightx2v_train.data import build_data +from lightx2v_train.data import build_data, build_sample_processor from lightx2v_train.infer import build_inferencer from lightx2v_train.model_zoo import build_model from lightx2v_train.runtime import cleanup_distributed, init_distributed, load_config, setup_logger @@ -75,9 +75,11 @@ def _build_low_model_for_dual_infer(config, reference_model): low_config["model"].update(copy.deepcopy(low_override)) low_model = build_model(low_config) low_model.load_components( - transformer_only=True, - reference_model=reference_model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + low_model.reuse_frozen_components_from(reference_model) _load_full_checkpoint_for_infer( low_model, low_config["model"], @@ -93,8 +95,13 @@ def main(): setup_logger(config) try: + sample_processor = build_sample_processor(config) model = build_model(config) - model.load_components() + model.load_components( + load_transformer=True, + load_vae=True, + load_condition_encoder=True, + ) _load_full_checkpoint_for_infer( model, config.get("model", {}), @@ -106,7 +113,11 @@ def main(): model.load_lora_for_infer(lora_path) apply_fsdp2(model, config) - dataloader_val = build_data(config, train_or_val="val") + dataloader_val = build_data( + config, + train_or_val="val", + sample_processor=sample_processor, + ) inferencer = build_inferencer(config) inferencer.set_model(model) diff --git a/lightx2v_train/lightx2v_train/data/__init__.py b/lightx2v_train/lightx2v_train/data/__init__.py index 040473df7..5d29a2020 100644 --- a/lightx2v_train/lightx2v_train/data/__init__.py +++ b/lightx2v_train/lightx2v_train/data/__init__.py @@ -1,4 +1,4 @@ -from lightx2v_train.utils.registry import build_data +from lightx2v_train.utils.registry import build_data, build_sample_processor def __getattr__(name): @@ -6,31 +6,25 @@ def __getattr__(name): from .image_dataset import build_image_dataset return build_image_dataset - if name == "build_libero_fastwam_dataset": - from .libero.dataset import build_libero_fastwam_dataset + if name == "build_training_cache_dataset": + from .training_cache_dataset import build_training_cache_dataset - return build_libero_fastwam_dataset - if name in {"build_latent_dataset", "build_prompt_dataset", "build_video_dataset"}: - from .video_dataset import build_latent_dataset, build_prompt_dataset, build_video_dataset + return build_training_cache_dataset + if name in {"build_prompt_dataset", "build_video_dataset"}: + from .video_dataset import build_prompt_dataset, build_video_dataset return { - "build_latent_dataset": build_latent_dataset, "build_prompt_dataset": build_prompt_dataset, "build_video_dataset": build_video_dataset, }[name] - if name == "prepare_data": - from .preparation import prepare_data - - return prepare_data raise AttributeError(f"module {__name__!r} has no attribute {name!r}") __all__ = [ "build_data", + "build_sample_processor", "build_image_dataset", - "build_libero_fastwam_dataset", - "build_latent_dataset", "build_prompt_dataset", + "build_training_cache_dataset", "build_video_dataset", - "prepare_data", ] diff --git a/lightx2v_train/lightx2v_train/data/image_dataset.py b/lightx2v_train/lightx2v_train/data/image_dataset.py index 8d362d32e..8b5f85fe4 100644 --- a/lightx2v_train/lightx2v_train/data/image_dataset.py +++ b/lightx2v_train/lightx2v_train/data/image_dataset.py @@ -1,10 +1,7 @@ import json -import math import random from pathlib import Path -import numpy as np -import torch from PIL import Image from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler @@ -12,48 +9,54 @@ from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size from lightx2v_train.utils.registry import DATA_REGISTER -_BICUBIC = getattr(Image, "Resampling", Image).BICUBIC - - -def _resize_to_target_area(image, target_area): - width, height = image.size - ratio = width / height - target_width = round(math.sqrt(target_area * ratio) / 16) * 16 - target_height = round(math.sqrt(target_area / ratio) / 16) * 16 - scale = max(target_width / width, target_height / height) - scaled_width = round(width * scale) - scaled_height = round(height * scale) - image = image.resize((scaled_width, scaled_height), resample=_BICUBIC) - left = (scaled_width - target_width) // 2 - top = (scaled_height - target_height) // 2 - return image.crop((left, top, left + target_width, top + target_height)) - class ImageDataset(Dataset): + """Decode image records and run an injected sample processor.""" + def __init__( self, metadata_paths, - target_area=1024 * 1024, prompt_dropout_rate=0.0, + sample_processor=None, + preserve_records=False, + unconditional_prompt=" ", ): - self.target_area = target_area self.prompt_dropout_rate = prompt_dropout_rate + self.sample_processor = sample_processor + self.preserve_records = preserve_records + self.unconditional_prompt = unconditional_prompt self.samples = [] for path in metadata_paths: path = Path(path) self.samples.extend(self._load_metadata_samples(path, data_dir=path.parent)) if not self.samples: raise ValueError(f"No valid image samples found in {metadata_paths}") + if self.sample_processor is None: + raise ValueError("Raw image records require a sample_processor.") def __len__(self): return len(self.samples) + def cache_source_record(self, index): + record = self.samples[index] + if "_original_record" not in record: + raise RuntimeError("ImageDataset was not configured with preserve_records=true.") + return dict(record["_original_record"]) + + def cache_source_prompt(self, index): + return self.samples[index]["prompt"] + def __getitem__(self, index): record = self.samples[index] prompt = record["prompt"] if random.random() < self.prompt_dropout_rate: - prompt = " " + prompt = self.unconditional_prompt + sample = self.load_sample(record, prompt=prompt) + if self.preserve_records: + sample["meta"]["dataset_index"] = index + return sample + def load_sample(self, record, prompt=None): inputs = {} if record.get("target_image") is not None: inputs["target_image"] = self.load_image(record["target_image"]) @@ -69,7 +72,12 @@ def __getitem__(self, index): meta["target_height"] = int(record["target_height"]) if record.get("target_width") is not None: meta["target_width"] = int(record["target_width"]) - return {"inputs": inputs, "conditioning": {"prompt": prompt}, "meta": meta} + sample = { + "inputs": inputs, + "conditioning": {"prompt": record["prompt"] if prompt is None else prompt}, + "meta": meta, + } + return sample if self.sample_processor is None else self.sample_processor(sample) def _load_metadata_samples(self, metadata_path, data_dir): if metadata_path.suffix != ".jsonl": @@ -89,14 +97,16 @@ def _normalize_record(self, record, data_dir): raise ValueError("Each metadata record must include prompt.") source_images = record.get("source_images", []) - - return { + normalized = { "target_image": self._resolve_path(target_image, data_dir) if target_image is not None else None, "prompt": str(prompt).strip(), "source_images": [self._resolve_path(p, data_dir) for p in source_images], "target_height": record.get("target_height"), "target_width": record.get("target_width"), } + if self.preserve_records: + normalized["_original_record"] = dict(record) + return normalized def _resolve_path(self, path, data_dir): path = Path(path) @@ -104,33 +114,45 @@ def _resolve_path(self, path, data_dir): return path return data_dir / path - def load_image(self, image_path): - image = Image.open(image_path).convert("RGB") - image = _resize_to_target_area(image, self.target_area) - return torch.from_numpy(np.asarray(image).astype(np.float32) / 127.5 - 1.0).permute(2, 0, 1) + @staticmethod + def load_image(image_path): + with Image.open(image_path) as image: + return image.convert("RGB").copy() @DATA_REGISTER("image_dataset") -def build_image_dataset(data_config_split, train_or_val="train"): +def build_image_dataset( + data_config_split, + train_or_val="train", + sample_processor=None, + unconditional_prompt=" ", +): data_path = data_config_split["data_path"] assert isinstance(data_path, list), f"config['data'][{train_or_val!r}]['data_path'] must be a list" - target_area = data_config_split.get("target_area", 1024 * 1024) prompt_dropout_rate = data_config_split.get("prompt_dropout_rate", 0.0) num_workers = data_config_split.get("num_workers", 8) shuffle = data_config_split.get("shuffle", train_or_val == "train") dataset = ImageDataset( metadata_paths=[Path(p) for p in data_path], - target_area=target_area, prompt_dropout_rate=prompt_dropout_rate, + sample_processor=sample_processor, + preserve_records=data_config_split.get("preserve_records", False), + unconditional_prompt=unconditional_prompt, ) dp_world_size = get_data_parallel_world_size() sampler = DistributedSampler(dataset, num_replicas=dp_world_size, rank=get_data_parallel_rank(), shuffle=shuffle) if dp_world_size > 1 and train_or_val == "train" else None + loader_kwargs = {} + if num_workers > 0: + loader_kwargs["persistent_workers"] = data_config_split.get("persistent_workers", False) + loader_kwargs["prefetch_factor"] = data_config_split.get("prefetch_factor", 2) return DataLoader( dataset, batch_size=1, shuffle=shuffle if sampler is None else False, sampler=sampler, num_workers=num_workers, + pin_memory=data_config_split.get("pin_memory", False), + **loader_kwargs, ) diff --git a/lightx2v_train/lightx2v_train/data/libero/__init__.py b/lightx2v_train/lightx2v_train/data/libero/__init__.py deleted file mode 100644 index 41a716905..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .dataset import build_libero_fastwam_dataset - -__all__ = ["build_libero_fastwam_dataset"] diff --git a/lightx2v_train/lightx2v_train/data/libero/dataset.py b/lightx2v_train/lightx2v_train/data/libero/dataset.py deleted file mode 100644 index 5e20b265e..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/dataset.py +++ /dev/null @@ -1,222 +0,0 @@ -import json -from copy import deepcopy -from pathlib import Path - -import torch -from loguru import logger -from torch.utils.data import DataLoader, DistributedSampler - -from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size -from lightx2v_train.utils.registry import DATA_REGISTER - -from .processor import FastWAMProcessor -from .robot_video_dataset import RobotVideoDataset - - -def _default_shape_meta(): - return { - "images": [ - {"key": "image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]}, - {"key": "wrist_image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]}, - ], - "action": [{"key": "default", "raw_shape": 7, "shape": 7}], - "state": [{"key": "default", "raw_shape": 8, "shape": 8}], - } - - -def _dataset_roots(config): - values = config.get("dataset_dirs") - if isinstance(values, (str, Path)): - values = [values] - if not values: - raise ValueError("LIBERO data.dataset_dirs must contain at least one dataset.") - return [Path(value).expanduser().resolve() for value in values] - - -def _load_dataset_features(roots): - result = [] - for root in roots: - info_path = root / "meta" / "info.json" - if not info_path.is_file(): - raise FileNotFoundError(f"LIBERO dataset metadata does not exist: {info_path}") - with info_path.open(encoding="utf-8") as handle: - info = json.load(handle) - features = info.get("features") - if not isinstance(features, dict): - raise ValueError(f"LIBERO dataset metadata has no feature map: {info_path}") - result.append((root, features)) - return result - - -def _feature_width(feature_sets, feature_key): - widths = {} - for root, features in feature_sets: - if feature_key not in features: - raise ValueError(f"LIBERO dataset {root} is missing feature: {feature_key}") - shape = features[feature_key].get("shape") - if not isinstance(shape, list) or len(shape) != 1: - raise ValueError(f"LIBERO feature {feature_key} must have a one-dimensional shape in {root}, got {shape}") - widths[root] = int(shape[0]) - unique_widths = set(widths.values()) - if len(unique_widths) != 1: - raise ValueError(f"LIBERO feature {feature_key} has inconsistent widths: {widths}") - return unique_widths.pop() - - -def _image_channels(feature_sets, image_key): - feature_key = f"observation.images.{image_key}" - channels = {} - for root, features in feature_sets: - if feature_key not in features: - raise ValueError(f"LIBERO dataset {root} is missing camera feature: {feature_key}") - feature = features[feature_key] - shape = feature.get("shape") - if feature.get("dtype") != "video" or not isinstance(shape, list) or len(shape) != 3: - raise ValueError(f"LIBERO camera feature {feature_key} must be a 3D video in {root}, got {feature}") - channels[root] = int(feature.get("info", {}).get("video.channels", shape[-1])) - unique_channels = set(channels.values()) - if len(unique_channels) != 1: - raise ValueError(f"LIBERO camera {image_key} has inconsistent channel counts: {channels}") - return unique_channels.pop() - - -def _normalize_vector_meta(configured, group, raw_width): - items = [{"key": "default"}] if configured is None else configured - if not isinstance(items, list) or not items: - raise ValueError(f"shape_meta.{group} must be a non-empty list when provided.") - - normalized = [] - for value in items: - item = {"key": value} if isinstance(value, str) else deepcopy(value) - if not isinstance(item, dict) or not item.get("key"): - raise ValueError(f"Invalid shape_meta.{group} entry: {value}") - if item["key"] != "default": - raise ValueError(f"LIBERO currently supports only shape_meta.{group} key 'default', got {item['key']!r}") - configured_raw_width = int(item.get("raw_shape", raw_width)) - configured_width = int(item.get("shape", configured_raw_width)) - if configured_raw_width != raw_width: - raise ValueError(f"shape_meta.{group}.default raw width does not match dataset metadata: configured={configured_raw_width}, dataset={raw_width}") - if configured_width != configured_raw_width: - raise ValueError(f"shape_meta.{group}.default does not support changing dimensions: raw={configured_raw_width}, processed={configured_width}") - normalized.append({"key": item["key"], "raw_shape": configured_raw_width, "shape": configured_width}) - return normalized - - -def _resolve_shape_meta(config): - roots = _dataset_roots(config) - feature_sets = _load_dataset_features(roots) - configured = deepcopy(config.get("shape_meta") or {}) - if not isinstance(configured, dict): - raise ValueError(f"shape_meta must be a mapping, got {type(configured).__name__}") - - image_size = config.get("image_size", (224, 224)) - if not isinstance(image_size, (list, tuple)) or len(image_size) != 2: - raise ValueError(f"data.image_size must be [height, width], got {image_size}") - image_size = [int(value) for value in image_size] - if any(value <= 0 for value in image_size): - raise ValueError(f"data.image_size values must be positive, got {image_size}") - - configured_images = configured.get("images") - if configured_images is None: - configured_images = _default_shape_meta()["images"] - if not isinstance(configured_images, list) or not configured_images: - raise ValueError("shape_meta.images must be a non-empty list.") - - images = [] - for value in configured_images: - item = {"key": value} if isinstance(value, str) else deepcopy(value) - if not isinstance(item, dict) or not item.get("key"): - raise ValueError(f"Invalid shape_meta.images entry: {value}") - channels = _image_channels(feature_sets, item["key"]) - target_shape = [int(dimension) for dimension in item.get("shape", [channels, *image_size])] - if target_shape != [channels, *image_size]: - raise ValueError(f"shape_meta.images.{item['key']} shape must match data.image_size and dataset channels: configured={target_shape}, expected={[channels, *image_size]}") - images.append({"key": item["key"], "shape": target_shape}) - - return { - "images": images, - "action": _normalize_vector_meta(configured.get("action"), "action", _feature_width(feature_sets, "action")), - "state": _normalize_vector_meta( - configured.get("state"), - "state", - _feature_width(feature_sets, "observation.state"), - ), - } - - -class DatasetSliceRepeat(torch.utils.data.Dataset): - def __init__(self, dataset, max_samples=None, dataset_repeat=1): - self.dataset = dataset - self.base_len = len(dataset) if max_samples is None else min(int(max_samples), len(dataset)) - self.dataset_repeat = max(1, int(dataset_repeat)) - if self.base_len <= 0: - raise RuntimeError("LIBERO dataset is empty after applying max_samples.") - - def __len__(self): - return self.base_len * self.dataset_repeat - - def __getitem__(self, index): - return self.dataset[index % self.base_len] - - -def _path(value): - return str(Path(value).expanduser().resolve()) - - -def _build_dataset(config, split): - shape_meta = _resolve_shape_meta(config) - num_frames = int(config.get("num_frames", 33)) - processor = FastWAMProcessor(shape_meta, num_frames) - dataset_dirs = _dataset_roots(config) - - dataset = RobotVideoDataset( - dataset_dirs=[str(item) for item in dataset_dirs], - shape_meta=shape_meta, - processor=processor, - text_embedding_cache_dir=_path(config["text_embedding_cache_dir"]), - pretrained_norm_stats=_path(config["pretrained_norm_stats"]), - num_frames=num_frames, - context_len=int(config.get("context_len", 128)), - val_set_proportion=float(config.get("val_set_proportion", 0.0)), - is_training_set=bool(config.get("is_training_set", split == "train")), - global_sample_stride=int(config.get("global_sample_stride", 1)), - action_video_freq_ratio=int(config.get("action_video_freq_ratio", 4)), - skip_padding_as_possible=bool(config.get("skip_padding_as_possible", False)), - max_padding_retry=int(config.get("max_padding_retry", 3)), - video_backend=config.get("video_backend"), - ) - logger.info("[data] built LIBERO FastWAM {} dataset size={}", split, len(dataset)) - return DatasetSliceRepeat( - dataset, - max_samples=config.get("max_samples"), - dataset_repeat=config.get("dataset_repeat", 1), - ) - - -@DATA_REGISTER("libero_fastwam_dataset") -def build_libero_fastwam_dataset(data_config, train_or_val="train"): - dataset = _build_dataset(data_config, train_or_val) - if data_config.get("return_dataset", False): - return dataset - - sampler = None - shuffle = bool(data_config.get("shuffle", train_or_val == "train")) - world_size = get_data_parallel_world_size() - if train_or_val == "train" and world_size > 1: - sampler = DistributedSampler( - dataset, - num_replicas=world_size, - rank=get_data_parallel_rank(), - shuffle=shuffle, - drop_last=bool(data_config.get("drop_last", False)), - ) - shuffle = False - return DataLoader( - dataset, - batch_size=int(data_config.get("batch_size", 1)), - shuffle=shuffle, - sampler=sampler, - num_workers=int(data_config.get("num_workers", 4)), - pin_memory=bool(data_config.get("pin_memory", torch.cuda.is_available())), - drop_last=bool(data_config.get("drop_last", False)), - ) diff --git a/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py b/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py deleted file mode 100644 index a213441af..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py +++ /dev/null @@ -1,187 +0,0 @@ -import bisect -import json -from dataclasses import dataclass -from functools import lru_cache -from pathlib import Path - -import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq -import torch -from torch.utils.data import Dataset - -from .video_decoder import decode_video_frames - - -@dataclass(frozen=True) -class Episode: - root: Path - index: int - length: int - data_path: str - video_path: str - chunks_size: int - tasks: dict[int, str] - - def parquet_path(self) -> Path: - return self.root / self.data_path.format( - episode_chunk=self.index // self.chunks_size, - episode_index=self.index, - ) - - def camera_path(self, key: str) -> Path: - return self.root / self.video_path.format( - episode_chunk=self.index // self.chunks_size, - episode_index=self.index, - video_key=key, - ) - - -def _read_json(path: Path): - with path.open(encoding="utf-8") as handle: - return json.load(handle) - - -def _read_jsonl(path: Path): - with path.open(encoding="utf-8") as handle: - return [json.loads(line) for line in handle if line.strip()] - - -def _column_to_tensor(column: pa.ChunkedArray) -> torch.Tensor: - column = column.combine_chunks() - if pa.types.is_list(column.type) or pa.types.is_fixed_size_list(column.type): - array = np.asarray(column.to_pylist()) - else: - array = column.to_numpy(zero_copy_only=False) - return torch.from_numpy(np.array(array, copy=True)) - - -class LiberoLeRobotDataset(Dataset): - """Minimal local reader for FastWAM's processed LIBERO LeRobot v2.1 data.""" - - def __init__( - self, - dataset_dirs, - image_keys, - state_key, - action_key, - num_frames, - global_sample_stride=1, - val_set_proportion=0.0, - is_training_set=True, - seed=42, - video_backend=None, - ): - if num_frames < 2: - raise ValueError(f"num_frames must be at least 2, got {num_frames}") - self.image_keys = list(image_keys) - self.state_key = state_key - self.action_key = action_key - self.num_frames = int(num_frames) - self.global_sample_stride = int(global_sample_stride) - self.video_backend = video_backend - - episodes = [] - fps_values = set() - for dataset_dir in dataset_dirs: - root = Path(dataset_dir).expanduser().resolve() - info = _read_json(root / "meta" / "info.json") - tasks = {int(item["task_index"]): item["task"] for item in _read_jsonl(root / "meta" / "tasks.jsonl")} - episode_meta = _read_jsonl(root / "meta" / "episodes.jsonl") - fps_values.add(int(info["fps"])) - self._validate_features(root, info["features"]) - - indices = list(range(len(episode_meta))) - if val_set_proportion >= 1e-6: - rng = np.random.default_rng(seed) - rng.shuffle(indices) - split = int(len(indices) * (1.0 - val_set_proportion)) - indices = indices[:split] if is_training_set else indices[split:] - - by_index = {int(item["episode_index"]): item for item in episode_meta} - for episode_index in indices: - item = by_index[episode_index] - episodes.append( - Episode( - root=root, - index=episode_index, - length=int(item["length"]), - data_path=info["data_path"], - video_path=info["video_path"], - chunks_size=int(info["chunks_size"]), - tasks=tasks, - ) - ) - - if not episodes: - raise RuntimeError("No LIBERO episodes were selected.") - if len(fps_values) != 1: - raise ValueError(f"All LIBERO datasets must have the same fps, got {sorted(fps_values)}") - self.fps = fps_values.pop() - self.episodes = episodes - self._episode_ends = np.cumsum([episode.length for episode in episodes]).tolist() - - def _validate_features(self, root, features): - required = [*self.image_keys, self.state_key, self.action_key, "timestamp", "task_index"] - missing = [key for key in required if key not in features] - if missing: - raise ValueError(f"LIBERO dataset {root} is missing features: {missing}") - non_video = [key for key in self.image_keys if features[key].get("dtype") != "video"] - if non_video: - raise ValueError(f"LIBERO camera features must be videos, got: {non_video}") - - def __len__(self): - return self._episode_ends[-1] - - @lru_cache(maxsize=32) - def _load_episode(self, episode_position): - episode = self.episodes[episode_position] - table = pq.read_table(episode.parquet_path()) - return {key: _column_to_tensor(table[key]) for key in table.column_names} - - def _locate(self, index): - if index < 0: - index += len(self) - if index < 0 or index >= len(self): - raise IndexError(f"Index {index} out of bounds for dataset of size {len(self)}") - episode_position = bisect.bisect_right(self._episode_ends, index) - episode_start = 0 if episode_position == 0 else self._episode_ends[episode_position - 1] - return episode_position, index - episode_start - - def _window(self, values, frame_index, length): - offsets = torch.arange(length, dtype=torch.long) * self.global_sample_stride - indices = frame_index + offsets - is_pad = indices >= values.shape[0] - indices.clamp_(max=values.shape[0] - 1) - return values[indices], is_pad, indices - - def __getitem__(self, index): - episode_position, frame_index = self._locate(index) - episode = self.episodes[episode_position] - data = self._load_episode(episode_position) - - state, state_is_pad, observation_indices = self._window(data[self.state_key], frame_index, self.num_frames) - action, action_is_pad, _ = self._window(data[self.action_key], frame_index, self.num_frames - 1) - timestamps = data["timestamp"][observation_indices].float().tolist() - tolerance = max(1e-4, 1.0 / self.fps - 1e-4) - images = { - key: decode_video_frames( - episode.camera_path(key), - timestamps, - tolerance_s=tolerance, - backend=self.video_backend, - ) - for key in self.image_keys - } - - task_index = int(data["task_index"][frame_index].item()) - return { - "idx": int(index), - "task": episode.tasks[task_index], - "action": {"default": action.float()}, - "state": {"default": state.float()}, - "images": {key.removeprefix("observation.images."): value for key, value in images.items()}, - "action_is_pad": action_is_pad, - "state_is_pad": state_is_pad, - "image_is_pad": state_is_pad.clone(), - } diff --git a/lightx2v_train/lightx2v_train/data/libero/normalizer.py b/lightx2v_train/lightx2v_train/data/libero/normalizer.py deleted file mode 100644 index 91cbbac46..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/normalizer.py +++ /dev/null @@ -1,69 +0,0 @@ -import json - -import torch - - -def _to_tensor_tree(value): - if isinstance(value, dict): - return {key: _to_tensor_tree(item) for key, item in value.items()} - if isinstance(value, list): - try: - return torch.tensor(value, dtype=torch.float32) - except (TypeError, ValueError): - return [_to_tensor_tree(item) for item in value] - return value - - -def load_dataset_stats(path): - with open(path, encoding="utf-8") as handle: - return _to_tensor_tree(json.load(handle)) - - -class FieldNormalizer: - def __init__(self, stats, mode="min/max"): - if mode == "z-score": - scale = 1.0 / (stats["std"] + 1e-8) - offset = -stats["mean"] * scale - else: - low_key, high_key = ("min", "max") if mode == "min/max" else ("q01", "q99") - lower = stats[low_key] - upper = stats[high_key] - width = upper - lower - constant = width < 1e-4 - width = width.clone() - width[constant] = 2.0 - scale = 2.0 / width - offset = -1.0 - scale * lower - offset[constant] = -lower[constant] - self.scale = scale - self.offset = offset - - def forward(self, value): - return (value * self.scale + self.offset).clamp_(-5.0, 5.0) - - def backward(self, value): - return (value - self.offset) / self.scale - - -class LinearNormalizer: - def __init__(self, shape_meta, stats, mode="min/max"): - self.normalizers = {"action": {}, "state": {}} - for group in self.normalizers: - for item in shape_meta[group]: - key = item["key"] - global_stats = {name.removeprefix("global_"): value for name, value in stats[group][key].items() if name.startswith("global_")} - self.normalizers[group][key] = FieldNormalizer(global_stats, mode=mode) - - def forward(self, batch): - for group, fields in self.normalizers.items(): - if group not in batch: - continue - for key, normalizer in fields.items(): - batch[group][key] = normalizer.forward(batch[group][key]) - return batch - - def backward(self, batch): - for group, fields in self.normalizers.items(): - for key, normalizer in fields.items(): - batch[group][key] = normalizer.backward(batch[group][key]) - return batch diff --git a/lightx2v_train/lightx2v_train/data/libero/preparation.py b/lightx2v_train/lightx2v_train/data/libero/preparation.py deleted file mode 100644 index 38524228c..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/preparation.py +++ /dev/null @@ -1,307 +0,0 @@ -import gc -import hashlib -import json -import os -import shutil -from concurrent.futures import ThreadPoolExecutor -from pathlib import Path - -import torch -from loguru import logger - -from lightx2v_train.runtime.distributed import barrier, is_main_process - -from .dataset import _resolve_shape_meta -from .lerobot_dataset import LiberoLeRobotDataset -from .robot_video_dataset import PROMPT_TEMPLATE - -TEXT_ENCODER_ID = "wan22ti2v5b" -TEXT_EMBEDDING_DIM = 4096 - - -def _resolve_path(value): - return Path(value).expanduser().resolve() - - -def _dataset_dirs(split_config): - values = split_config.get("dataset_dirs") - if isinstance(values, (str, Path)): - values = [values] - return [_resolve_path(value) for value in values] - - -def _atomic_json_dump(payload, path): - path.parent.mkdir(parents=True, exist_ok=True) - temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") - with temporary_path.open("w", encoding="utf-8") as handle: - json.dump(_to_serializable(payload), handle, ensure_ascii=True, indent=2) - os.replace(temporary_path, path) - - -def _atomic_torch_save(payload, path): - path.parent.mkdir(parents=True, exist_ok=True) - temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") - torch.save(payload, temporary_path) - os.replace(temporary_path, path) - - -def _to_serializable(value): - if isinstance(value, torch.Tensor): - return value.detach().cpu().tolist() - if isinstance(value, dict): - return {key: _to_serializable(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_to_serializable(item) for item in value] - return value - - -def _sliding_window_with_replication(value, window_size): - offsets = torch.arange(window_size, dtype=torch.long) - indices = torch.arange(value.shape[0], dtype=torch.long)[:, None] + offsets[None, :] - indices.clamp_(max=value.shape[0] - 1) - return value[indices] - - -def _episode_field_stats(value): - return { - "min": value.amin(0), - "max": value.amax(0), - "q01": torch.quantile(value, 0.01, dim=0, keepdim=False), - "q99": torch.quantile(value, 0.99, dim=0, keepdim=False), - "mean": value.mean(0), - "var": value.var(0), - } - - -def _aggregate_field_stats(episode_stats): - values = {name: torch.stack([stats[name] for stats in episode_stats]) for name in ("min", "max", "q01", "q99", "mean", "var")} - stepwise_mean = values["mean"].mean(0) - stepwise_std = (values["var"] + (values["mean"] - stepwise_mean) ** 2).mean(0).sqrt() - global_mean = values["mean"].mean((0, 1)) - global_std = (values["var"] + (values["mean"] - global_mean) ** 2).mean((0, 1)).sqrt() - stepwise_min = values["min"].amin(0) - stepwise_max = values["max"].amax(0) - stepwise_q01 = values["q01"].amin(0) - stepwise_q99 = values["q99"].amax(0) - return { - "stepwise_min": stepwise_min, - "stepwise_max": stepwise_max, - "global_min": stepwise_min.amin(0), - "global_max": stepwise_max.amax(0), - "stepwise_q01": stepwise_q01, - "stepwise_q99": stepwise_q99, - "global_q01": stepwise_q01.amin(0), - "global_q99": stepwise_q99.amax(0), - "stepwise_mean": stepwise_mean, - "stepwise_std": stepwise_std, - "global_mean": global_mean, - "global_std": global_std, - } - - -def calculate_dataset_stats(split_config): - shape_meta = _resolve_shape_meta(split_config) - if len(shape_meta["action"]) != 1 or len(shape_meta["state"]) != 1: - raise ValueError("Automatic LIBERO stats currently require one action field and one state field.") - - image_keys = [f"observation.images.{item['key']}" for item in shape_meta["images"]] - dataset = LiberoLeRobotDataset( - dataset_dirs=[str(path) for path in _dataset_dirs(split_config)], - image_keys=image_keys, - state_key="observation.state", - action_key="action", - num_frames=int(split_config.get("num_frames", 33)), - global_sample_stride=int(split_config.get("global_sample_stride", 1)), - val_set_proportion=float(split_config.get("val_set_proportion", 0.0)), - is_training_set=True, - video_backend=split_config.get("video_backend"), - ) - action_horizon = int(split_config.get("num_frames", 33)) - 1 - - def process_episode(episode_position): - raw = dataset._load_episode(episode_position) - state = raw[dataset.state_key].float().unsqueeze(1) - action = _sliding_window_with_replication(raw[dataset.action_key].float(), action_horizon) - return _episode_field_stats(state), _episode_field_stats(action) - - worker_count = min(16, len(dataset.episodes), os.cpu_count() or 1) - with ThreadPoolExecutor(max_workers=worker_count) as executor: - episode_results = list(executor.map(process_episode, range(len(dataset.episodes)))) - - state_key = shape_meta["state"][0]["key"] - action_key = shape_meta["action"][0]["key"] - return { - "state": {state_key: _aggregate_field_stats([result[0] for result in episode_results])}, - "action": {action_key: _aggregate_field_stats([result[1] for result in episode_results])}, - "num_episodes": len(dataset.episodes), - "num_transition": len(dataset), - } - - -def _collect_prompts(split_configs): - prompts = [] - seen = set() - for split_config in split_configs: - for dataset_dir in _dataset_dirs(split_config): - tasks_path = dataset_dir / "meta" / "tasks.jsonl" - if not tasks_path.is_file(): - raise FileNotFoundError(f"LIBERO tasks file does not exist: {tasks_path}") - with tasks_path.open(encoding="utf-8") as handle: - for line in handle: - if not line.strip(): - continue - task = str(json.loads(line)["task"]) - prompt = PROMPT_TEMPLATE.format(task=task) - if prompt not in seen: - seen.add(prompt) - prompts.append(prompt) - return prompts - - -def _text_cache_path(cache_dir, prompt, context_len): - digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest() - return cache_dir / f"{digest}.t5_len{context_len}.{TEXT_ENCODER_ID}.pt" - - -def _valid_text_cache(path, context_len): - if not path.is_file(): - return False - try: - payload = torch.load(path, map_location="cpu", weights_only=True) - return tuple(payload["context"].shape) == (context_len, TEXT_EMBEDDING_DIM) and tuple(payload["mask"].shape) == (context_len,) - except Exception: - return False - - -def _validate_text_caches(split_configs): - for split_config in split_configs: - cache_dir = _resolve_path(split_config["text_embedding_cache_dir"]) - context_len = int(split_config.get("context_len", 128)) - prompts = _collect_prompts([split_config]) - missing = [_text_cache_path(cache_dir, prompt, context_len) for prompt in prompts if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len)] - if missing: - raise FileNotFoundError(f"LIBERO text embedding cache is incomplete: missing={len(missing)}/{len(prompts)} first={missing[0]}") - - -def precompute_text_embeddings(model_path, cache_dir, context_len, prompts): - missing_prompts = [prompt for prompt in prompts if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len)] - if not missing_prompts: - logger.info("[data-preflight] text embeddings already cached: prompts={} path={}", len(prompts), cache_dir) - return - if not torch.cuda.is_available(): - raise RuntimeError("Automatic FastWAM text embedding generation requires CUDA.") - - from lightx2v_train.model_zoo.native.wan.modules.t5 import T5EncoderModel - - model_root = _resolve_path(model_path) - checkpoint_path = model_root / "models_t5_umt5-xxl-enc-bf16.pth" - tokenizer_path = model_root / "google" / "umt5-xxl" - if not checkpoint_path.is_file(): - raise FileNotFoundError(f"Wan text encoder checkpoint does not exist: {checkpoint_path}") - if not tokenizer_path.is_dir(): - raise FileNotFoundError(f"Wan tokenizer directory does not exist: {tokenizer_path}") - - device = torch.device("cuda", torch.cuda.current_device()) - logger.info( - "[data-preflight] generating text embeddings: missing={}/{} context_len={} path={}", - len(missing_prompts), - len(prompts), - context_len, - cache_dir, - ) - with torch.random.fork_rng(devices=[torch.cuda.current_device()]): - encoder = T5EncoderModel( - text_len=context_len, - dtype=torch.bfloat16, - device=device, - checkpoint_path=str(checkpoint_path), - tokenizer_path=str(tokenizer_path), - ) - encoder.model.eval() - batch_size = 16 - with torch.no_grad(): - for start in range(0, len(missing_prompts), batch_size): - batch_prompts = missing_prompts[start : start + batch_size] - input_ids, mask = encoder.tokenizer(batch_prompts, return_mask=True, add_special_tokens=True) - input_ids = input_ids.to(device) - mask = mask.to(device=device, dtype=torch.bool) - context = encoder.model(input_ids, mask) - for index, prompt in enumerate(batch_prompts): - _atomic_torch_save( - { - "context": context[index].detach().to(device="cpu", dtype=torch.bfloat16).contiguous(), - "mask": mask[index].detach().to(device="cpu", dtype=torch.bool).contiguous(), - }, - _text_cache_path(cache_dir, prompt, context_len), - ) - - del context, input_ids, mask, encoder - gc.collect() - torch.cuda.empty_cache() - logger.info("[data-preflight] generated {} text embeddings in {}", len(missing_prompts), cache_dir) - - -def _is_missing(value): - return value is None or not str(value).strip() - - -def prepare_libero_fastwam_assets(config): - data_config = config["data"] - train_config = data_config["train"] - val_config = data_config.get("val") - output_dir = _resolve_path(config["training"]["output_dir"]) - - configured_stats = train_config.get("pretrained_norm_stats") - if _is_missing(configured_stats): - stats_path = output_dir / "dataset_stats.json" - if is_main_process() and not stats_path.is_file(): - logger.info("[data-preflight] calculating LIBERO normalization statistics") - stats = calculate_dataset_stats(train_config) - _atomic_json_dump(stats, stats_path) - logger.info( - "[data-preflight] saved normalization statistics: episodes={} transitions={} path={}", - stats["num_episodes"], - stats["num_transition"], - stats_path, - ) - barrier() - else: - stats_path = _resolve_path(configured_stats) - if not stats_path.is_file(): - raise FileNotFoundError(f"Configured LIBERO normalization stats do not exist: {stats_path}") - eval_stats_path = output_dir / "dataset_stats.json" - if is_main_process() and stats_path != eval_stats_path: - eval_stats_path.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(stats_path, eval_stats_path) - barrier() - - train_config["pretrained_norm_stats"] = str(stats_path) - if val_config is not None and _is_missing(val_config.get("pretrained_norm_stats")): - val_config["pretrained_norm_stats"] = str(stats_path) - - configured_cache = train_config.get("text_embedding_cache_dir") - auto_text_cache = _is_missing(configured_cache) - cache_dir = output_dir / "text_embeds_cache" if auto_text_cache else _resolve_path(configured_cache) - train_config["text_embedding_cache_dir"] = str(cache_dir) - if val_config is not None and _is_missing(val_config.get("text_embedding_cache_dir")): - val_config["text_embedding_cache_dir"] = str(cache_dir) - - if auto_text_cache: - generated_splits = [train_config] - if val_config is not None and _resolve_path(val_config["text_embedding_cache_dir"]) == cache_dir: - generated_splits.append(val_config) - context_lengths = {int(split.get("context_len", 128)) for split in generated_splits} - if len(context_lengths) != 1: - raise ValueError(f"Automatic text cache generation requires one context_len, got {sorted(context_lengths)}") - context_len = context_lengths.pop() - prompts = _collect_prompts(generated_splits) - if is_main_process(): - precompute_text_embeddings( - model_path=config["model"]["model_path"], - cache_dir=cache_dir, - context_len=context_len, - prompts=prompts, - ) - barrier() - - _validate_text_caches([split for split in (train_config, val_config) if split is not None]) diff --git a/lightx2v_train/lightx2v_train/data/libero/processor.py b/lightx2v_train/lightx2v_train/data/libero/processor.py deleted file mode 100644 index cd2747ba9..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/processor.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -import torchvision.transforms.functional as vision - -from .normalizer import LinearNormalizer - - -class ConcatLeftAlign: - def __init__(self, shape_meta): - self.action_meta = shape_meta["action"] - self.state_meta = shape_meta["state"] - - @staticmethod - def _merge(values, metadata): - return torch.cat([values[item["key"]] for item in metadata], dim=-1) - - @staticmethod - def _split(value, metadata): - result = {} - offset = 0 - for item in metadata: - width = int(item["shape"]) - result[item["key"]] = value[..., offset : offset + width] - offset += width - return result - - def forward(self, batch): - if "action" in batch: - batch["action"] = self._merge(batch["action"], self.action_meta) - batch["action_dim_is_pad"] = torch.zeros(batch["action"].shape[-1], dtype=torch.bool) - batch["state"] = self._merge(batch["state"], self.state_meta) - batch["state_dim_is_pad"] = torch.zeros(batch["state"].shape[-1], dtype=torch.bool) - return batch - - def backward(self, batch): - batch["action"] = self._split(batch["action"], self.action_meta) - batch["state"] = self._split(batch["state"], self.state_meta) - return batch - - -class FastWAMProcessor: - def __init__( - self, - shape_meta, - num_obs_steps, - image_size=(224, 224), - delta_action_dim_mask=(True, True, True, True, True, True, False), - ): - self.shape_meta = shape_meta - self.num_obs_steps = int(num_obs_steps) - self.image_size = tuple(image_size) - self.action_state_merger = ConcatLeftAlign(shape_meta) - self.delta_action_dim_mask = torch.as_tensor(delta_action_dim_mask, dtype=torch.bool) - self._normalizer = None - - @property - def normalizer(self): - if self._normalizer is None: - raise RuntimeError("FastWAM processor normalization statistics have not been loaded.") - return self._normalizer - - def set_normalizer_from_stats(self, stats): - self._normalizer = LinearNormalizer(self.shape_meta, stats) - - def _validate_action_state(self, data): - for group in ("action", "state"): - for item in self.shape_meta[group]: - actual = data[group][item["key"]].shape[-1] - expected = int(item["raw_shape"]) - if actual != expected: - raise ValueError(f"{group}.{item['key']} width mismatch: expected {expected}, got {actual}") - - def _process_images(self, images): - processed = [] - for item in self.shape_meta["images"]: - image = images[item["key"]] - if image.ndim != 4: - raise ValueError(f"Image sequence must be [T,C,H,W], got {tuple(image.shape)}") - image = vision.resize( - image.to(torch.float32), - self.image_size, - interpolation=vision.InterpolationMode.BILINEAR, - antialias=True, - ) - expected = (self.num_obs_steps, *tuple(item["shape"])) - if tuple(image.shape) != expected: - raise ValueError(f"Processed image shape mismatch: expected {expected}, got {tuple(image.shape)}") - processed.append(image) - return torch.stack(processed) - - def preprocess(self, data): - self._validate_action_state(data) - action_is_pad = torch.as_tensor(data["action_is_pad"], dtype=torch.bool) - action = data["action"]["default"] - if bool(action_is_pad.any()): - mask = action_is_pad[:, None] & self.delta_action_dim_mask[None, :] - action[mask] = 0.0 - - transformed = self.normalizer.forward({"action": data["action"], "state": data["state"]}) - transformed = self.action_state_merger.forward(transformed) - return { - "idx": data["idx"], - "instruction": str(data["task"]), - "pixel_values": self._process_images(data["images"]), - "image_is_pad": data["image_is_pad"], - "action": transformed["action"], - "action_is_pad": action_is_pad, - "action_dim_is_pad": transformed["action_dim_is_pad"], - "proprio": transformed["state"], - "proprio_is_pad": data["state_is_pad"], - "proprio_dim_is_pad": transformed["state_dim_is_pad"], - } diff --git a/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py b/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py deleted file mode 100644 index 76ff3d147..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py +++ /dev/null @@ -1,125 +0,0 @@ -import hashlib -import os -import traceback - -import numpy as np -import torch -from torch.utils.data import Dataset - -from .lerobot_dataset import LiberoLeRobotDataset -from .normalizer import load_dataset_stats - -PROMPT_TEMPLATE = "A video recorded from a robot's point of view executing the following instruction: {task}" - - -class RobotVideoDataset(Dataset): - def __init__( - self, - dataset_dirs, - shape_meta, - processor, - text_embedding_cache_dir, - pretrained_norm_stats, - num_frames=33, - context_len=128, - val_set_proportion=0.0, - is_training_set=True, - global_sample_stride=1, - action_video_freq_ratio=4, - skip_padding_as_possible=False, - max_padding_retry=3, - video_backend=None, - ): - image_keys = [f"observation.images.{item['key']}" for item in shape_meta["images"]] - self.lerobot_dataset = LiberoLeRobotDataset( - dataset_dirs=dataset_dirs, - image_keys=image_keys, - state_key="observation.state", - action_key="action", - num_frames=num_frames, - global_sample_stride=global_sample_stride, - val_set_proportion=val_set_proportion, - is_training_set=is_training_set, - video_backend=video_backend, - ) - self.num_frames = int(num_frames) - self.action_video_freq_ratio = int(action_video_freq_ratio) - if (self.num_frames - 1) % self.action_video_freq_ratio: - raise ValueError("num_frames - 1 must be divisible by action_video_freq_ratio") - if ((self.num_frames - 1) // self.action_video_freq_ratio) % 4: - raise ValueError("The number of future video frames must be divisible by the VAE temporal factor 4") - self.video_sample_indices = list(range(0, self.num_frames, self.action_video_freq_ratio)) - self.context_len = int(context_len) - self.text_embedding_cache_dir = os.path.abspath(os.path.expanduser(text_embedding_cache_dir)) - self.skip_padding_as_possible = bool(skip_padding_as_possible) - self.max_padding_retry = int(max_padding_retry) - - stats_path = os.path.abspath(os.path.expanduser(pretrained_norm_stats)) - if not os.path.isfile(stats_path): - raise FileNotFoundError(f"LIBERO normalization stats do not exist: {stats_path}") - processor.set_normalizer_from_stats(load_dataset_stats(stats_path)) - self.processor = processor - - def __len__(self): - return len(self.lerobot_dataset) - - def _sample_without_padding(self, index): - for attempt in range(self.max_padding_retry + 1): - sample = self.processor.preprocess(self.lerobot_dataset[index]) - if not self.skip_padding_as_possible: - return sample - has_padding = any(bool(sample[key].any()) for key in ("action_is_pad", "image_is_pad", "proprio_is_pad")) - if not has_padding or attempt == self.max_padding_retry: - return sample - index = np.random.randint(len(self)) - raise AssertionError("unreachable") - - def _cached_context(self, prompt): - digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest() - path = os.path.join( - self.text_embedding_cache_dir, - f"{digest}.t5_len{self.context_len}.wan22ti2v5b.pt", - ) - if not os.path.isfile(path): - raise FileNotFoundError(f"Missing cached text embedding: {path}") - payload = torch.load(path, map_location="cpu", weights_only=True) - context = payload["context"] - mask = payload["mask"].bool() - if context.ndim != 2 or context.shape[0] != self.context_len: - raise ValueError(f"Invalid cached context shape {tuple(context.shape)} in {path}") - if mask.shape != (self.context_len,): - raise ValueError(f"Invalid cached context mask shape {tuple(mask.shape)} in {path}") - context[~mask] = 0.0 - return context, torch.ones_like(mask) - - def _get(self, index): - sample = self._sample_without_padding(index) - video = sample["pixel_values"][:, self.video_sample_indices] - if video.shape[0] != 2: - raise ValueError(f"FastWAM LIBERO expects two cameras, got {video.shape[0]}") - video = torch.cat((video[0], video[1]), dim=-1) - video = video.mul(2.0).sub(1.0).permute(1, 0, 2, 3).contiguous() - - action = sample["action"] - proprio = sample["proprio"][:-1] - prompt = PROMPT_TEMPLATE.format(task=sample["instruction"]) - context, context_mask = self._cached_context(prompt) - return { - "video": video, - "action": action, - "proprio": proprio, - "prompt": prompt, - "context": context, - "context_mask": context_mask, - "image_is_pad": sample["image_is_pad"][self.video_sample_indices], - "action_is_pad": sample["action_is_pad"], - "proprio_is_pad": sample["proprio_is_pad"], - } - - def __getitem__(self, index): - try: - return self._get(index) - except Exception: - traceback.print_exc() - fallback = np.random.randint(len(self)) - return self._get(fallback) diff --git a/lightx2v_train/lightx2v_train/data/libero/video_decoder.py b/lightx2v_train/lightx2v_train/data/libero/video_decoder.py deleted file mode 100644 index 748be766b..000000000 --- a/lightx2v_train/lightx2v_train/data/libero/video_decoder.py +++ /dev/null @@ -1,67 +0,0 @@ -import importlib.util -import logging -import warnings -from pathlib import Path - -import torch -import torchvision - - -def _default_backend(): - if importlib.util.find_spec("torchcodec") is not None: - return "torchcodec" - return "pyav" - - -def decode_video_frames(video_path, timestamps, tolerance_s, backend=None): - backend = backend or _default_backend() - if backend == "torchcodec": - try: - return _decode_torchcodec(video_path, timestamps, tolerance_s) - except Exception as error: - warnings.warn(f"torchcodec decode failed ({type(error).__name__}: {error}); falling back to pyav.") - backend = "pyav" - if backend not in {"pyav", "video_reader"}: - raise ValueError(f"Unsupported video backend: {backend}") - return _decode_torchvision(video_path, timestamps, tolerance_s, backend) - - -def _decode_torchvision(video_path, timestamps, tolerance_s, backend): - torchvision.set_video_backend(backend) - reader = torchvision.io.VideoReader(str(video_path), "video") - reader.seek(min(timestamps), keyframes_only=backend == "pyav") - frames = [] - frame_timestamps = [] - for frame in reader: - frames.append(frame["data"]) - frame_timestamps.append(frame["pts"]) - if frame["pts"] >= max(timestamps): - break - if backend == "pyav": - reader.container.close() - if not frames: - raise RuntimeError(f"No frames decoded from {video_path}") - return _select_frames(frames, frame_timestamps, timestamps, tolerance_s, video_path) - - -def _decode_torchcodec(video_path, timestamps, tolerance_s): - from torchcodec.decoders import VideoDecoder - - decoder = VideoDecoder(str(video_path), device="cpu", seek_mode="approximate") - fps = decoder.metadata.average_fps - batch = decoder.get_frames_at(indices=[round(timestamp * fps) for timestamp in timestamps]) - frames = list(batch.data) - frame_timestamps = batch.pts_seconds.tolist() - return _select_frames(frames, frame_timestamps, timestamps, tolerance_s, video_path) - - -def _select_frames(frames, frame_timestamps, query_timestamps, tolerance_s, video_path): - query = torch.as_tensor(query_timestamps, dtype=torch.float32) - available = torch.as_tensor(frame_timestamps, dtype=torch.float32) - distances = torch.cdist(query[:, None], available[:, None], p=1) - minimum, indices = distances.min(dim=1) - if not bool((minimum <= tolerance_s).all()): - logging.warning("Video timestamp mismatch for %s: max error %.6fs", video_path, minimum.max()) - raise RuntimeError(f"Unable to decode requested timestamps from {Path(video_path)}") - selected = torch.stack([frames[index] for index in indices.tolist()]) - return selected.to(torch.float32).div_(255.0) diff --git a/lightx2v_train/lightx2v_train/data/preparation.py b/lightx2v_train/lightx2v_train/data/preparation.py deleted file mode 100644 index fc3bb0ea9..000000000 --- a/lightx2v_train/lightx2v_train/data/preparation.py +++ /dev/null @@ -1,7 +0,0 @@ -def prepare_data(config): - data_config = config.get("data", {}) - data_names = {split_config.get("name") for split_config in data_config.values() if isinstance(split_config, dict)} - if "libero_fastwam_dataset" in data_names: - from .libero.preparation import prepare_libero_fastwam_assets - - prepare_libero_fastwam_assets(config) diff --git a/lightx2v_train/lightx2v_train/data/training_cache.py b/lightx2v_train/lightx2v_train/data/training_cache.py new file mode 100644 index 000000000..9149bebd8 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/training_cache.py @@ -0,0 +1,61 @@ +import hashlib +import json + +CACHE_SCHEMA_VERSION = 2 + +_DMD_METHODS = {"dmd", "autoregressive_dmd", "phased_dmd", "sgmd"} +_RUNTIME_MODEL_KEYS = { + "cache_condition_encoder_cpu_offload", + "cache_encoder_cpu_offload", + "load_condition_encoder", + "load_text_encoder", + "load_transformer", + "load_vae", +} + +_RUNTIME_DATA_KEYS = { + "data_path", + "dataset_repeat", + "drop_last", + "max_samples", + "num_workers", + "persistent_workers", + "pin_memory", + "prefetch_factor", + "preserve_records", + "shuffle", +} + + +def preserve_cache_dtype(key): + return isinstance(key, str) and (key.endswith("_ids") or key.endswith("_mask")) + + +def training_cache_info(config): + model = config["model"] + data = config.get("data", {}) + train_data = data.get("train", {}) + training = config["training"] + method = training["method"] + objective_key = "dmd" if method in _DMD_METHODS else method + signature_data = { + "model": {key: value for key, value in model.items() if key not in _RUNTIME_MODEL_KEYS}, + "data": { + "processor": data.get("processor", {}), + "train": {key: value for key, value in train_data.items() if key not in _RUNTIME_DATA_KEYS}, + }, + "training": { + "method": method, + "objective": training.get(objective_key, {}), + "teacher": training.get("teacher", {}), + }, + "target_latent_mode": "mode", + } + serialized = json.dumps(signature_data, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str) + return { + "schema_version": CACHE_SCHEMA_VERSION, + "model_name": model["name"], + "model_path": model["pretrained_model_name_or_path"], + "training_method": method, + "signature": hashlib.sha256(serialized.encode("utf-8")).hexdigest(), + } diff --git a/lightx2v_train/lightx2v_train/data/training_cache_dataset.py b/lightx2v_train/lightx2v_train/data/training_cache_dataset.py new file mode 100644 index 000000000..f53fce1de --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/training_cache_dataset.py @@ -0,0 +1,154 @@ +"""Model-agnostic dataset for caches produced by ``cache_data.py``.""" + +import json +import random +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, Dataset +from torch.utils.data.distributed import DistributedSampler + +from lightx2v_train.data.training_cache import CACHE_SCHEMA_VERSION +from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size +from lightx2v_train.utils.registry import DATA_REGISTER + + +class TrainingCacheDataset(Dataset): + """Load the common ``cache_data.jsonl`` format for every model family.""" + + def __init__( + self, + metadata_paths, + prompt_dropout_rate=0.0, + unconditional_prompt=" ", + expected_cache_info=None, + ): + self.prompt_dropout_rate = float(prompt_dropout_rate) + self.unconditional_prompt = unconditional_prompt + self.expected_cache_info = expected_cache_info + self.samples = [] + for metadata_path in metadata_paths: + self.samples.extend(self._read_manifest(Path(metadata_path))) + if not self.samples: + raise ValueError(f"No training-cache records found in {metadata_paths}.") + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + record = self.samples[index] + cache_path = record["training_cache"] + cache = torch.load(cache_path, map_location="cpu", weights_only=True) + self._validate_cache(cache, cache_path, record["prompt"]) + + conditioning = cache["conditioning"] + use_unconditional = random.random() < self.prompt_dropout_rate + if use_unconditional and "unconditional" not in conditioning: + raise ValueError(f"Training cache {cache_path} has no unconditional condition for prompt dropout.") + conditioning["active"] = "unconditional" if use_unconditional else "positive" + conditioning["prompt"] = conditioning.get("unconditional_prompt", self.unconditional_prompt) if use_unconditional else record["prompt"] + cache["meta"]["training_cache_path"] = str(cache_path) + return cache + + def _read_manifest(self, metadata_path): + if metadata_path.suffix.lower() != ".jsonl": + raise ValueError(f"Training-cache manifests must be .jsonl files, got: {metadata_path}") + records = [] + with metadata_path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + record = json.loads(line) + prompt = record.get("prompt") + cache_path = record.get("training_cache") + if prompt is None or cache_path is None or not str(cache_path).strip(): + raise ValueError(f"Training-cache record {metadata_path}:{line_number} must include prompt and training_cache.") + cache_path = Path(cache_path) + if not cache_path.is_absolute(): + cache_path = metadata_path.parent / cache_path + records.append( + { + "prompt": str(prompt), + "training_cache": cache_path, + } + ) + return records + + def _validate_cache(self, cache, path, prompt): + required = {"inputs", "conditioning", "meta"} + if not isinstance(cache, dict) or not required.issubset(cache): + raise ValueError(f"Invalid training cache at {path}: expected inputs, conditioning, and meta mappings.") + if not all(isinstance(cache[key], dict) for key in required): + raise ValueError(f"Invalid training cache at {path}: inputs, conditioning, and meta must be mappings.") + + cache_info = cache.get("cache_info") + if not isinstance(cache_info, dict): + raise ValueError(f"Invalid training cache metadata at {path}.") + if cache_info.get("schema_version") != CACHE_SCHEMA_VERSION: + raise ValueError(f"Unsupported training cache schema at {path}: {cache_info.get('schema_version')!r}.") + if self.expected_cache_info is not None: + for key, expected in self.expected_cache_info.items(): + if cache_info.get(key) != expected: + raise ValueError(f"Training cache {path} has incompatible {key}: expected {expected!r}, got {cache_info.get(key)!r}.") + + conditioning = cache["conditioning"] + if "positive" not in conditioning: + raise ValueError(f"Invalid training cache at {path}: conditioning.positive is missing.") + if conditioning.get("prompt") != prompt: + raise ValueError(f"Prompt in {path} does not match cache_data.jsonl. Rebuild the training cache.") + + +@DATA_REGISTER("training_cache_dataset") +def build_training_cache_dataset( + data_config_split, + train_or_val="train", + unconditional_prompt=" ", + expected_cache_info=None, +): + if train_or_val != "train": + raise ValueError("training_cache_dataset is only valid for the training split.") + data_paths = data_config_split["data_path"] + if isinstance(data_paths, (str, Path)): + data_paths = [data_paths] + + dataset = TrainingCacheDataset( + metadata_paths=data_paths, + prompt_dropout_rate=data_config_split.get("prompt_dropout_rate", 0.0), + unconditional_prompt=unconditional_prompt, + expected_cache_info=expected_cache_info, + ) + world_size = get_data_parallel_world_size() + shuffle = data_config_split.get("shuffle", True) + sampler = ( + DistributedSampler( + dataset, + num_replicas=world_size, + rank=get_data_parallel_rank(), + shuffle=shuffle, + drop_last=data_config_split.get("drop_last", False), + ) + if world_size > 1 + else None + ) + num_workers = int(data_config_split.get("num_workers", 8)) + loader_kwargs = {} + if num_workers > 0: + loader_kwargs["persistent_workers"] = data_config_split.get("persistent_workers", True) + loader_kwargs["prefetch_factor"] = data_config_split.get("prefetch_factor", 2) + return DataLoader( + dataset, + batch_size=1, + shuffle=shuffle if sampler is None else False, + sampler=sampler, + num_workers=num_workers, + pin_memory=data_config_split.get("pin_memory", True), + drop_last=data_config_split.get("drop_last", False), + collate_fn=_single_sample_collate, + **loader_kwargs, + ) + + +def _single_sample_collate(samples): + if len(samples) != 1: + raise ValueError(f"Cached training requires batch_size=1, got {len(samples)} samples.") + return samples[0] diff --git a/lightx2v_train/lightx2v_train/data/utils.py b/lightx2v_train/lightx2v_train/data/utils.py index 265b4b5d5..11d9420fd 100644 --- a/lightx2v_train/lightx2v_train/data/utils.py +++ b/lightx2v_train/lightx2v_train/data/utils.py @@ -3,6 +3,7 @@ import math import random import sys +from dataclasses import dataclass from pathlib import Path import imageio @@ -14,6 +15,18 @@ _BILINEAR = getattr(Image, "Resampling", Image).BILINEAR +@dataclass(frozen=True) +class VideoFrameSelection: + frame_ids: tuple[int, ...] + start_time: float + + +def require_singleton_dataloader(dataloader, name): + batch_size = getattr(dataloader, "batch_size", None) + if batch_size != 1: + raise ValueError(f"{name} must use DataLoader(batch_size=1); custom batch samplers and physical batch sizes greater than 1 are not supported, got {batch_size!r}.") + + def to_list(value): if value is None: return [] @@ -88,10 +101,6 @@ def read_records(path, prompt_column="prompt", prompt_index=0): "audio_path", "image", "image_path", - "video_latent_path", - "audio_latent_path", - "condition_path", - "negative_condition_path", } if set(header).intersection(known_columns): for row in reader: @@ -180,20 +189,31 @@ def raw_frame_id(self, sequence_id, raw_frame_rate, total_raw_frames): frame_id = int(round(target_time * raw_frame_rate)) return min(frame_id, total_raw_frames - 1) - def frame_ids(self, reader): - raw_frame_rate = reader.get_meta_data().get("fps", self.frame_rate) + def sample(self, reader): + raw_frame_rate = float(reader.get_meta_data().get("fps") or self.frame_rate) + if raw_frame_rate <= 0: + raise ValueError(f"Video frame rate must be positive, got {raw_frame_rate}.") total_raw_frames = int(reader.count_frames()) num_frames = self.sample_count(reader) max_start = max(0, self.available_frames(reader) - num_frames) start = random.randint(0, max_start) if self.random_start and max_start > 0 else 0 - return [self.raw_frame_id(start + frame_id, raw_frame_rate, total_raw_frames) for frame_id in range(num_frames)] + frame_ids = tuple(self.raw_frame_id(start + frame_id, raw_frame_rate, total_raw_frames) for frame_id in range(num_frames)) + sampling_frame_rate = self.frame_rate if self.fix_frame_rate else raw_frame_rate + return VideoFrameSelection( + frame_ids=frame_ids, + start_time=start / sampling_frame_rate, + ) + + def frame_ids(self, reader): + return list(self.sample(reader).frame_ids) -def load_video_tensor(video_path, height, width, frame_sampler): +def load_video_tensor(video_path, height, width, frame_sampler, *, return_start_time=False): reader = imageio.get_reader(video_path) try: + selection = frame_sampler.sample(reader) frames = [] - for frame_id in frame_sampler.frame_ids(reader): + for frame_id in selection.frame_ids: frame = Image.fromarray(reader.get_data(frame_id)).convert("RGB") frame_width, frame_height = frame.size scale = max(width / frame_width, height / frame_height) @@ -207,4 +227,5 @@ def load_video_tensor(video_path, height, width, frame_sampler): frames.append(torch.from_numpy(array).permute(2, 0, 1)) finally: reader.close() - return torch.stack(frames, dim=1) + video = torch.stack(frames, dim=1) + return (video, selection.start_time) if return_start_time else video diff --git a/lightx2v_train/lightx2v_train/data/video_dataset.py b/lightx2v_train/lightx2v_train/data/video_dataset.py index a092198cb..f1b38984c 100644 --- a/lightx2v_train/lightx2v_train/data/video_dataset.py +++ b/lightx2v_train/lightx2v_train/data/video_dataset.py @@ -1,9 +1,8 @@ -import io import json import random +from collections import defaultdict from pathlib import Path -import numpy as np import torch from loguru import logger from torch.utils.data import DataLoader, DistributedSampler @@ -18,71 +17,11 @@ to_list, ) from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size +from lightx2v_train.utils.image_size_buckets import parse_image_size_buckets from lightx2v_train.utils.registry import DATA_REGISTER METADATA_SUFFIXES = {".jsonl", ".json", ".csv"} PROMPT_SUFFIXES = {".txt", ".list"} -CONDITION_KEYS = ( - "prompt_embed", - "video_prompt_embeds", - "audio_prompt_embeds", - "prompt_embeds", - "text_token_tags", - "prompt_attention_mask", - "video_context", - "audio_context", - "context_mask", -) - - -def _strip_condition_batch(value): - if torch.is_tensor(value): - if value.ndim >= 2 and value.shape[0] == 1: - return value[0] - return value - if isinstance(value, dict): - return {key: _strip_condition_batch(item) for key, item in value.items()} - if isinstance(value, list): - return [_strip_condition_batch(item) for item in value] - if isinstance(value, tuple): - return tuple(_strip_condition_batch(item) for item in value) - return value - - -def _condition_payload(item): - if isinstance(item, (list, tuple)): - return _strip_condition_batch(item) - if not isinstance(item, dict): - raise TypeError(f"Condition payload must be a dict/list/tuple, got {type(item)!r}.") - if "conditioning" in item: - item = item["conditioning"] - if "positive" in item: - conditions = item["positive"] - elif "conditions" in item: - conditions = item["conditions"] - elif "condition" in item: - conditions = item["condition"] - else: - conditions = {key: item[key] for key in CONDITION_KEYS if key in item} - if conditions is None or (isinstance(conditions, (dict, list, tuple)) and not conditions): - raise KeyError("Condition payload must contain positive/conditions/condition or prompt embedding tensors.") - return _strip_condition_batch(conditions) - - -def _video_latent_payload(data): - if torch.is_tensor(data): - return {"latents": data} - if not isinstance(data, dict): - return data - normalized = dict(data) - latents = normalized.get("latents") - if not torch.is_tensor(latents) or latents.dim() != 2: - return normalized - num_frames = int(normalized["num_frames"]) - height = int(normalized["height"]) - width = int(normalized["width"]) - normalized["latents"] = latents.reshape(num_frames, height, width, latents.shape[-1]).permute(3, 0, 1, 2).contiguous() - return normalized def _metadata_path(path): @@ -97,20 +36,6 @@ def _metadata_path(path): return path -def _is_lmdb_path(path): - path = Path(path) - return path.is_dir() and ((path / "data.mdb").is_file() or (path / "lock.mdb").is_file()) - - -def _resolve_required_path(value, base_dir, key): - path = resolve_data_path(value, base_dir) - if path is None: - return None - if not path.is_file(): - raise FileNotFoundError(f"{key} points to a missing file: {path}") - return path - - class VideoDataset(torch.utils.data.Dataset): def __init__( self, @@ -134,6 +59,9 @@ def __init__( frame_rate=24, fix_frame_rate=False, decode_retries=3, + preserve_records=False, + sample_processor=None, + unconditional_prompt=" ", ): self.metadata_paths = [_metadata_path(path) for path in to_list(metadata_paths)] self.height = int(height) @@ -152,6 +80,9 @@ def __init__( self.skip_missing = bool(skip_missing) self.max_samples = max_samples self.decode_retries = max(1, int(decode_retries)) + self.preserve_records = bool(preserve_records) + self.sample_processor = sample_processor + self.unconditional_prompt = unconditional_prompt self.frame_sampler = VideoFrameSampler( num_frames=num_frames, frame_rate=frame_rate, @@ -186,6 +117,10 @@ def _load_samples(self): audio_value = record_value(row, self.audio_column, "audio_path", "audio") audio_path = resolve_data_path(audio_value, metadata_path.parent, self.audio_roots, subdirs=("audio", "audios")) + if getattr(self.sample_processor, "requires_audio", False) and (audio_path is None or not audio_path.is_file()): + if self.skip_missing: + continue + raise FileNotFoundError(f"Audio path points to a missing file: {audio_path or audio_value}") if audio_path is not None: meta["audio_path"] = str(audio_path) @@ -213,29 +148,44 @@ def _load_samples(self): prompt = " ".join(prompt_path.read_text(encoding="utf-8").split()) meta["prompt_path"] = str(prompt_path) - samples.append({"prompt": prompt, "meta": meta}) + sample = {"prompt": prompt, "meta": meta} + if self.preserve_records: + sample["_original_record"] = dict(row) + samples.append(sample) if self.max_samples is not None and len(samples) >= int(self.max_samples): return samples return samples def _load_video(self, video_path): - return load_video_tensor(video_path, self.height, self.width, self.frame_sampler) + return load_video_tensor( + video_path, + self.height, + self.width, + self.frame_sampler, + return_start_time=True, + ) def __getitem__(self, index): base_index = index % len(self.samples) last_error = None for retry_id in range(self.decode_retries): - record = self.samples[(base_index + retry_id) % len(self.samples)] + sample_index = (base_index + retry_id) % len(self.samples) + record = self.samples[sample_index] meta = dict(record["meta"]) try: prompt = record["prompt"] if random.random() < self.prompt_dropout_rate: - prompt = " " - return { - "inputs": {"video": self._load_video(meta["video_path"])}, + prompt = self.unconditional_prompt + video, video_start_time = self._load_video(meta["video_path"]) + meta["video_start_time"] = video_start_time + sample = { + "inputs": {"video": video}, "conditioning": {"prompt": prompt}, "meta": meta, } + if self.preserve_records: + sample["meta"]["dataset_index"] = sample_index + return sample if self.sample_processor is None else self.sample_processor(sample) except Exception as error: last_error = error logger.warning("Failed to load video {}: {}", meta.get("video_path"), error) @@ -244,6 +194,15 @@ def __getitem__(self, index): def __len__(self): return len(self.samples) * self.dataset_repeat + def cache_source_record(self, index): + record = self.samples[index] + if "_original_record" not in record: + raise RuntimeError("VideoDataset was not configured with preserve_records=true.") + return dict(record["_original_record"]) + + def cache_source_prompt(self, index): + return self.samples[index]["prompt"] + class PromptDataset(torch.utils.data.Dataset): def __init__( @@ -252,6 +211,8 @@ def __init__( dataset_repeat=1, max_samples=None, prompt_dropout_rate=0.0, + preserve_records=False, + unconditional_prompt=" ", ): self.prompt_paths = [Path(path) for path in to_list(prompt_paths)] if not self.prompt_paths: @@ -259,6 +220,8 @@ def __init__( self.dataset_repeat = int(dataset_repeat) self.max_samples = None if max_samples is None else int(max_samples) self.prompt_dropout_rate = float(prompt_dropout_rate) + self.preserve_records = bool(preserve_records) + self.unconditional_prompt = unconditional_prompt self.samples = self._load_prompts() if not self.samples: raise RuntimeError(f"No prompts found from data_path={prompt_paths}") @@ -279,23 +242,38 @@ def _load_prompts(self): negative_prompt = json.dumps(negative_prompt, ensure_ascii=False, separators=(",", ":")) else: negative_prompt = str(negative_prompt).strip() - samples.append( - { - "prompt": prompt, - "negative_prompt": negative_prompt, - "meta": {"prompt_path": str(path), "row_index": row_index}, - } - ) + target_height = record_value(record, "target_height") + target_width = record_value(record, "target_width") + if (target_height is None) != (target_width is None): + raise ValueError(f"Prompt metadata must provide target_height and target_width together: {path}:{row_index + 1}") + if target_height is not None: + target_height = int(target_height) + target_width = int(target_width) + if target_height <= 0 or target_width <= 0: + raise ValueError(f"Prompt target dimensions must be positive, got {target_height}x{target_width}: {path}:{row_index + 1}") + meta = {"prompt_path": str(path), "row_index": row_index} + if target_height is not None: + meta["target_height"] = target_height + meta["target_width"] = target_width + sample = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "meta": meta, + } + if self.preserve_records: + sample["_original_record"] = dict(record) + samples.append(sample) if self.max_samples is not None and len(samples) >= self.max_samples: return samples return samples def __getitem__(self, index): - sample = self.samples[index % len(self.samples)] + sample_index = index % len(self.samples) + sample = self.samples[sample_index] prompt = sample["prompt"] if random.random() < self.prompt_dropout_rate: - prompt = " " - return { + prompt = self.unconditional_prompt + result = { "inputs": {}, "conditioning": { "prompt": prompt, @@ -303,216 +281,184 @@ def __getitem__(self, index): }, "meta": dict(sample["meta"]), } + if self.preserve_records: + result["meta"]["dataset_index"] = sample_index + return result def __len__(self): return len(self.samples) * self.dataset_repeat + def cache_source_record(self, index): + sample = self.samples[index] + if "_original_record" not in sample: + raise RuntimeError("PromptDataset was not configured with preserve_records=true.") + return dict(sample["_original_record"]) -class LatentDataset(torch.utils.data.Dataset): - def __init__( - self, - data_paths, - dataset_repeat=1, - max_samples=None, - prompt_column="caption", - prompt_index=0, - negative_condition_path=None, - ): - self.paths = [Path(path) for path in to_list(data_paths)] - if not self.paths: - raise ValueError("latent_dataset requires data_path.") - self.dataset_repeat = int(dataset_repeat) - self.max_samples = None if max_samples is None else int(max_samples) - self.prompt_column = prompt_column - self.prompt_index = int(prompt_index) - self.samples = [] - self.lmdb_envs = [] - self.negative_condition = self._load_negative_condition(negative_condition_path) - - for path in self.paths: - self._index_path(path) - if self.max_samples is not None: - self.samples = self.samples[: self.max_samples] - if not self.samples: - raise RuntimeError(f"No usable latent samples found from data_path={data_paths}.") - logger.info("[data] latent_dataset samples={} repeat={}", len(self.samples), self.dataset_repeat) + def cache_source_prompt(self, index): + return self.samples[index]["prompt"] - def _index_path(self, path): - if _is_lmdb_path(path): - self._index_lmdb(path) - return - self._index_metadata(_metadata_path(path)) - - def _load_negative_condition(self, negative_condition_path): - if negative_condition_path is None: - for path in self.paths: - base = path if path.is_dir() else path.parent - candidate = base / "negative_condition.pt" - if candidate.is_file(): - negative_condition_path = candidate - break - if negative_condition_path is None: + def size_bucket(self, index): + sample = self.samples[index % len(self.samples)] + meta = sample["meta"] + height = meta.get("target_height") + width = meta.get("target_width") + if height is None or width is None: return None - return _condition_payload(torch.load(negative_condition_path, map_location="cpu", weights_only=False)) - - def _index_metadata(self, metadata_path): - for row in read_records(metadata_path, prompt_column=self.prompt_column, prompt_index=self.prompt_index): - self.samples.append({"type": "metadata", "row": row, "base_dir": str(metadata_path.parent)}) - - def _index_lmdb(self, data_path): - try: - import lmdb - except ImportError as error: - raise ImportError("latent_dataset LMDB input requires the 'lmdb' Python package.") from error - - env = lmdb.open(str(data_path), readonly=True, lock=False, readahead=False, meminit=False) - with env.begin() as txn: - format_bytes = txn.get(b"__format__") - sample_count_bytes = txn.get(b"sample_count") or txn.get(b"num_samples") - - env_index = len(self.lmdb_envs) - if format_bytes is not None and format_bytes.decode() == "sample_pt": - if sample_count_bytes is None: - raise KeyError("sample_pt LMDB dataset requires sample_count.") - sample_count = int(sample_count_bytes.decode()) - self.lmdb_envs.append({"env": env, "format": "sample_pt", "path": str(data_path)}) - for row_index in range(sample_count): - self.samples.append({"type": "lmdb_sample", "env_index": env_index, "row_index": row_index}) - return + return int(height), int(width) - latents_shape = self._get_lmdb_shape(env, "latents") - self.lmdb_envs.append({"env": env, "format": "wan_latents", "latents_shape": latents_shape, "path": str(data_path)}) - for row_index in range(latents_shape[0]): - self.samples.append({"type": "lmdb_wan", "env_index": env_index, "row_index": row_index}) - - @staticmethod - def _get_lmdb_shape(env, array_name): - with env.begin() as txn: - shape_bytes = txn.get(f"{array_name}_shape".encode()) - if shape_bytes is None: - raise KeyError(f"{array_name}_shape not found in LMDB dataset.") - return tuple(map(int, shape_bytes.decode().split())) - - @staticmethod - def _retrieve_lmdb_row(env, array_name, dtype, row_index, shape=None): - data_key = f"{array_name}_{row_index}_data".encode() - with env.begin() as txn: - row_bytes = txn.get(data_key) - if row_bytes is None: - raise KeyError(f"{data_key!r} not found in LMDB dataset.") - if dtype is str: - return row_bytes.decode() - array = np.frombuffer(row_bytes, dtype=dtype) - if shape is not None and len(shape) > 0: - array = array.reshape(shape) - return array - def __getitem__(self, index): - sample = self.samples[index % len(self.samples)] - if sample["type"] == "metadata": - return self._load_metadata_sample(sample["row"], Path(sample["base_dir"])) - if sample["type"] in {"lmdb_sample", "lmdb_wan"}: - return self._load_lmdb_sample(sample) - raise AssertionError(f"Unhandled latent dataset sample type: {sample['type']}") - - def _load_metadata_sample(self, row, base_dir): - inputs = {} - conditioning = {} - meta = {} - height = record_value(row, "target_height", "height") - width = record_value(row, "target_width", "width") - if height not in (None, "") and width not in (None, ""): - meta["target_height"] = int(height) - meta["target_width"] = int(width) - for key in ("id", "width", "height", "fps", "frames", "duration", "num_frames"): - if isinstance(row, dict) and key in row: - meta[key] = row[key] - for key in ("video", "video_path", "audio", "audio_path", "image", "image_path"): - value = record_value(row, key) - path = resolve_data_path(value, base_dir) - if path is not None: - meta[key] = str(path) - prompt = prompt_text(row, self.prompt_column, self.prompt_index) - if prompt: - conditioning["prompt"] = prompt - - video_latent_path = _resolve_required_path(record_value(row, "video_latent_path"), base_dir, "video_latent_path") - if video_latent_path is not None: - video_payload = _video_latent_payload(torch.load(video_latent_path, map_location="cpu", weights_only=True)) - inputs["video_latents"] = video_payload - if torch.is_tensor(video_payload): - inputs["latents"] = video_payload - elif isinstance(video_payload, dict) and torch.is_tensor(video_payload.get("latents")): - inputs["latents"] = video_payload["latents"] - meta["video_latent_path"] = str(video_latent_path) - - audio_latent_path = _resolve_required_path(record_value(row, "audio_latent_path"), base_dir, "audio_latent_path") - if audio_latent_path is not None: - inputs["audio_latents"] = torch.load(audio_latent_path, map_location="cpu", weights_only=True) - meta["audio_latent_path"] = str(audio_latent_path) - - condition_path = _resolve_required_path(record_value(row, "condition_path"), base_dir, "condition_path") - if condition_path is not None: - condition_item = torch.load(condition_path, map_location="cpu", weights_only=False) - positive = _condition_payload(condition_item) - conditioning["positive"] = positive - meta["condition_path"] = str(condition_path) - self._add_negative_condition(conditioning, condition_item) - elif self.negative_condition is not None: - conditioning["negative"] = self.negative_condition - - row_negative_path = _resolve_required_path(record_value(row, "negative_condition_path"), base_dir, "negative_condition_path") - if row_negative_path is not None: - conditioning["negative"] = _condition_payload(torch.load(row_negative_path, map_location="cpu", weights_only=False)) - meta["negative_condition_path"] = str(row_negative_path) - - if not inputs and "positive" not in conditioning: - raise ValueError(f"Latent metadata row must contain video_latent_path and/or condition_path: {row}") - return {"inputs": inputs, "conditioning": conditioning, "meta": meta} - - def _load_lmdb_sample(self, record): - source = self.lmdb_envs[record["env_index"]] - env = source["env"] - row_index = record["row_index"] - - if source.get("format") == "sample_pt": - data_key = f"sample_{row_index:08d}".encode() - with env.begin() as txn: - row_bytes = txn.get(data_key) - if row_bytes is None: - raise KeyError(f"{data_key!r} not found in LMDB dataset.") - sample = torch.load(io.BytesIO(row_bytes), map_location="cpu", weights_only=False) - sample = { - "inputs": sample.get("inputs", {}), - "conditioning": sample.get("conditioning", {}), - "meta": sample.get("meta", {}), - } - sample["meta"].setdefault("row_index", row_index) - sample["meta"].setdefault("lmdb_path", source.get("path")) - return sample - - latents_shape = source["latents_shape"] - latents = self._retrieve_lmdb_row(env, "latents", np.float16, row_index, shape=latents_shape[1:]) - if latents.ndim == 4: - latents = latents[None, ...] - latent_tchw = torch.tensor(latents, dtype=torch.float32)[-1] - prompt = self._retrieve_lmdb_row(env, "prompts", str, row_index) - return { - "inputs": {"latents": latent_tchw.permute(1, 0, 2, 3).contiguous()}, - "conditioning": {"prompt": prompt}, - "meta": {"row_index": row_index, "lmdb_path": source.get("path")}, - } +class SizeBucketSampler(torch.utils.data.Sampler): + """Keep all data-parallel ranks on one exact size bucket per step.""" - def _add_negative_condition(self, conditioning, item): - if self.negative_condition is not None: - conditioning["negative"] = self.negative_condition - elif isinstance(item, dict) and "negative" in item: - conditioning["negative"] = _condition_payload({"positive": item["negative"]}) - elif isinstance(item, dict) and "negative_conditions" in item: - conditioning["negative"] = _condition_payload({"positive": item["negative_conditions"]}) + def __init__( + self, + dataset, + num_replicas=1, + rank=0, + shuffle=True, + drop_last=False, + seed=0, + image_sizes=None, + ): + if not hasattr(dataset, "size_bucket"): + raise TypeError("Size-bucket sampling requires dataset.size_bucket(index).") + self.dataset = dataset + self.num_replicas = int(num_replicas) + self.rank = int(rank) + self.shuffle = bool(shuffle) + self.drop_last = bool(drop_last) + self.seed = int(seed) + self.epoch = 0 + self.image_size_buckets = parse_image_size_buckets(image_sizes) + if self.num_replicas <= 0: + raise ValueError(f"num_replicas must be positive, got {self.num_replicas}.") + if self.rank < 0 or self.rank >= self.num_replicas: + raise ValueError(f"rank must be in [0, {self.num_replicas}), got {self.rank}.") + + self.buckets = defaultdict(list) + for index in range(len(dataset)): + bucket = dataset.size_bucket(index) + if bucket is None: + raise ValueError(f"Size-bucket sampling requires target_height and target_width for every prompt sample; missing at dataset index {index}.") + self.buckets[tuple(bucket)].append(index) + self._validate_configured_buckets() + self.weighted = bool(self.image_size_buckets and self.image_size_buckets[0].ratio is not None) + self.num_steps = sum(self._bucket_total(len(indices)) // self.num_replicas for indices in self.buckets.values()) + self.num_samples = self.num_steps + if self.num_samples == 0: + raise ValueError("Size-bucket sampling produced no steps. Disable drop_last or add more samples to each size bucket.") + + def _bucket_total(self, size): + multiple = self.num_replicas + if self.drop_last: + return size // multiple * multiple + return (size + multiple - 1) // multiple * multiple + + def _validate_configured_buckets(self): + if not self.image_size_buckets: + return + configured = {bucket.spatial_size for bucket in self.image_size_buckets} + present = set(self.buckets) + unexpected = sorted(present - configured) + missing = sorted(configured - present) + if unexpected: + values = ", ".join(f"{height}x{width}" for height, width in unexpected) + raise ValueError(f"Prompt dataset contains target sizes not listed in training.dmd.image_sizes: {values}.") + if missing: + values = ", ".join(f"{height}x{width}" for height, width in missing) + raise ValueError(f"training.dmd.image_sizes contains buckets with no prompt samples: {values}.") + + def set_epoch(self, epoch): + self.epoch = int(epoch) + + def __iter__(self): + generator = torch.Generator() + generator.manual_seed(self.seed + self.epoch) + if self.weighted: + return iter(self._weighted_indices(generator)) + return iter(self._natural_indices(generator)) + + def _natural_indices(self, generator): + local_indices = [] + for bucket in sorted(self.buckets): + indices = list(self.buckets[bucket]) + if self.shuffle: + order = torch.randperm(len(indices), generator=generator).tolist() + indices = [indices[position] for position in order] + + total = self._bucket_total(len(indices)) + if total == 0: + continue + if total > len(indices): + repeat = (total + len(indices) - 1) // len(indices) + indices = (indices * repeat)[:total] + else: + indices = indices[:total] + + for start in range(0, total, self.num_replicas): + replica_group = indices[start : start + self.num_replicas] + local_indices.append(replica_group[self.rank]) + + if self.shuffle and len(local_indices) > 1: + order = torch.randperm(len(local_indices), generator=generator).tolist() + local_indices = [local_indices[position] for position in order] + return local_indices + + def _weighted_indices(self, generator): + bucket_order = [bucket.spatial_size for bucket in self.image_size_buckets] + weights = torch.tensor( + [bucket.ratio for bucket in self.image_size_buckets], + dtype=torch.double, + ) + exact_counts = weights / weights.sum() * self.num_steps + step_counts_tensor = exact_counts.floor().to(dtype=torch.int64) + remainder = self.num_steps - int(step_counts_tensor.sum().item()) + if remainder: + fractions = exact_counts - step_counts_tensor + selected = torch.multinomial( + fractions, + remainder, + replacement=False, + generator=generator, + ) + step_counts_tensor[selected] += 1 + + selected_buckets = [] + for bucket, step_count in zip(bucket_order, step_counts_tensor.tolist()): + selected_buckets.extend([bucket] * step_count) + if self.shuffle and len(selected_buckets) > 1: + order = torch.randperm(len(selected_buckets), generator=generator).tolist() + selected_buckets = [selected_buckets[position] for position in order] + + step_counts = defaultdict(int) + for bucket in selected_buckets: + step_counts[bucket] += 1 + + sampled_indices = {} + for bucket, step_count in step_counts.items(): + needed = step_count * self.num_replicas + source = self.buckets[bucket] + sampled = [] + while len(sampled) < needed: + if self.shuffle: + order = torch.randperm(len(source), generator=generator).tolist() + cycle = [source[position] for position in order] + else: + cycle = list(source) + sampled.extend(cycle[: needed - len(sampled)]) + sampled_indices[bucket] = sampled + + offsets = defaultdict(int) + local_indices = [] + for bucket in selected_buckets: + start = offsets[bucket] + replica_group = sampled_indices[bucket][start : start + self.num_replicas] + offsets[bucket] += self.num_replicas + local_indices.append(replica_group[self.rank]) + return local_indices def __len__(self): - return len(self.samples) * self.dataset_repeat + return self.num_samples def _build_dataloader(dataset, data_config, train_or_val): @@ -520,7 +466,23 @@ def _build_dataloader(dataset, data_config, train_or_val): sampler = None shuffle = data_config.get("shuffle", train_or_val == "train") drop_last = data_config.get("drop_last", False) - if train_or_val == "train" and dp_world_size > 1: + image_size_buckets = parse_image_size_buckets(data_config.get("image_sizes")) + if train_or_val == "train" and image_size_buckets and image_size_buckets[0].ratio is not None and not data_config.get("bucket_by_size", False): + raise ValueError("training.dmd.image_sizes ratio requires data.train.bucket_by_size=true.") + if train_or_val == "train" and data_config.get("bucket_by_size", False): + sampler = SizeBucketSampler( + dataset, + num_replicas=dp_world_size, + rank=get_data_parallel_rank(), + shuffle=shuffle, + drop_last=drop_last, + seed=data_config.get("bucket_seed", 0), + image_sizes=data_config.get("image_sizes"), + ) + shuffle = False + # SizeBucketSampler already pads or truncates every distributed step. + drop_last = False + elif train_or_val == "train" and dp_world_size > 1: sampler = DistributedSampler( dataset, num_replicas=dp_world_size, @@ -532,7 +494,7 @@ def _build_dataloader(dataset, data_config, train_or_val): return DataLoader( dataset, - batch_size=data_config.get("batch_size", 1), + batch_size=1, shuffle=shuffle if sampler is None else False, sampler=sampler, num_workers=data_config.get("num_workers", 8), @@ -542,7 +504,7 @@ def _build_dataloader(dataset, data_config, train_or_val): @DATA_REGISTER("video_dataset") -def build_video_dataset(data_config, train_or_val="train"): +def build_video_dataset(data_config, train_or_val="train", sample_processor=None): dataset = VideoDataset( metadata_paths=data_config["data_path"], height=data_config.get("height", 480), @@ -564,6 +526,9 @@ def build_video_dataset(data_config, train_or_val="train"): frame_rate=data_config.get("frame_rate", 24), fix_frame_rate=data_config.get("fix_frame_rate", False), decode_retries=data_config.get("decode_retries", 3), + preserve_records=data_config.get("preserve_records", False), + sample_processor=sample_processor, + unconditional_prompt=getattr(sample_processor, "unconditional_prompt", " "), ) return _build_dataloader(dataset, data_config, train_or_val) @@ -575,18 +540,6 @@ def build_prompt_dataset(data_config, train_or_val="train"): dataset_repeat=data_config.get("dataset_repeat", 1), max_samples=data_config.get("max_samples"), prompt_dropout_rate=data_config.get("prompt_dropout_rate", 0.0), - ) - return _build_dataloader(dataset, data_config, train_or_val) - - -@DATA_REGISTER("latent_dataset") -def build_latent_dataset(data_config, train_or_val="train"): - dataset = LatentDataset( - data_paths=data_config["data_path"], - dataset_repeat=data_config.get("dataset_repeat", 1), - max_samples=data_config.get("max_samples"), - prompt_column=data_config.get("prompt_column", "caption"), - prompt_index=data_config.get("prompt_index", 0), - negative_condition_path=data_config.get("negative_condition_path"), + preserve_records=data_config.get("preserve_records", False), ) return _build_dataloader(dataset, data_config, train_or_val) diff --git a/lightx2v_train/lightx2v_train/infer/__init__.py b/lightx2v_train/lightx2v_train/infer/__init__.py index e77173de7..5523ae5e2 100644 --- a/lightx2v_train/lightx2v_train/infer/__init__.py +++ b/lightx2v_train/lightx2v_train/infer/__init__.py @@ -4,7 +4,6 @@ _LAZY_EXPORTS = { "ImageInferencer": (".image", "ImageInferencer"), - "NativeImageInferencer": (".image_native", "NativeImageInferencer"), "LingBotVideoT2VInferencer": ( ".video", "LingBotVideoT2VInferencer", @@ -31,7 +30,6 @@ def __getattr__(name): __all__ = [ "build_inferencer", "ImageInferencer", - "NativeImageInferencer", "LingBotVideoT2VInferencer", "WanT2VDualInferencer", "WanT2VInferencer", diff --git a/lightx2v_train/lightx2v_train/infer/base.py b/lightx2v_train/lightx2v_train/infer/base.py index 48ba12503..099062290 100644 --- a/lightx2v_train/lightx2v_train/infer/base.py +++ b/lightx2v_train/lightx2v_train/infer/base.py @@ -45,10 +45,6 @@ def cfg_guided_denoise( timestep_or_sigma, pos_cond, ) - pred_pos = model.postprocess_denoiser_output( - pred_pos, - denoiser_input, - ) if self.enable_cfg: pred_neg = model.denoise( @@ -56,14 +52,10 @@ def cfg_guided_denoise( timestep_or_sigma, neg_cond, ) - pred_neg = model.postprocess_denoiser_output( - pred_neg, - denoiser_input, - ) - pred = pred_neg + self.guidance_scale * (pred_pos - pred_neg) + pred = model.apply_cfg(pred_pos, pred_neg, self.guidance_scale) else: pred = pred_pos - return pred + return model.postprocess_denoiser_output(pred, denoiser_input) @torch.no_grad() def infer(self): diff --git a/lightx2v_train/lightx2v_train/infer/dopsd_trajectory_viz.py b/lightx2v_train/lightx2v_train/infer/dopsd_trajectory_viz.py deleted file mode 100644 index 03dfe78e4..000000000 --- a/lightx2v_train/lightx2v_train/infer/dopsd_trajectory_viz.py +++ /dev/null @@ -1,47 +0,0 @@ -from PIL import Image, ImageDraw, ImageFont - - -def _fit_height(image, height): - if image.height == height: - return image - width = max(1, int(image.width * height / image.height)) - return image.resize((width, height), Image.Resampling.LANCZOS) - - -def save_student_teacher_trajectory_grid(student_step_images, teacher_step_images, save_path): - if len(student_step_images) != len(teacher_step_images): - raise ValueError("student and teacher trajectory lengths must match") - - pad = 12 - header_h = 32 - row_label_w = 56 - font = ImageFont.load_default() - num_steps = len(student_step_images) - if num_steps == 0: - return - - row_h = max(img.height for img in student_step_images + teacher_step_images) - student_cols = [_fit_height(img.convert("RGB"), row_h) for img in student_step_images] - teacher_cols = [_fit_height(img.convert("RGB"), row_h) for img in teacher_step_images] - - panel_w = student_cols[0].width + pad + teacher_cols[0].width - canvas_h = header_h + num_steps * (row_h + pad) + pad - canvas_w = row_label_w + pad + panel_w + pad - canvas = Image.new("RGB", (canvas_w, canvas_h), (255, 255, 255)) - draw = ImageDraw.Draw(canvas) - - header_y = 8 - draw.text((row_label_w + pad + 8, header_y), "Student", fill=(0, 0, 0), font=font) - draw.text((row_label_w + pad + student_cols[0].width + pad + 8, header_y), "Teacher", fill=(0, 0, 0), font=font) - - y = header_h - for step_idx, (student_img, teacher_img) in enumerate(zip(student_cols, teacher_cols)): - draw.text((8, y + (row_h - 10) // 2), f"t{step_idx}", fill=(0, 0, 0), font=font) - x_student = row_label_w + pad - canvas.paste(student_img, (x_student, y)) - x_teacher = x_student + student_img.width + pad - canvas.paste(teacher_img, (x_teacher, y)) - y += row_h + pad - - save_path = str(save_path) - canvas.save(save_path) diff --git a/lightx2v_train/lightx2v_train/infer/image.py b/lightx2v_train/lightx2v_train/infer/image.py index 284aae6a4..eb75483ff 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -13,20 +13,25 @@ def _record_has_source_images(record): return bool(record.get("source_images")) -def _target_hw_for_sample(sample, default_height, default_width): - h = sample["meta"].get("target_height") - w = sample["meta"].get("target_width") - if h is not None and w is not None: - return int(h), int(w) - source_images = sample["inputs"].get("source_images") - if source_images: - source_image = source_images[0] - return int(source_image.shape[-2]), int(source_image.shape[-1]) - return default_height, default_width - - @INFERENCER_REGISTER("image_infer") class ImageInferencer(BaseInferencer): + def _inference_sigmas(self, num_inference_steps, *, latent_hw=None): + sigmas = self.infer_config.get("sigmas") + pcm_solver_steps = self.infer_config.get("pcm_solver_steps") + if sigmas is not None and pcm_solver_steps is not None: + raise ValueError("Set only one of inference.sigmas and inference.pcm_solver_steps.") + if pcm_solver_steps is None: + return sigmas + + from lightx2v_train.trainers.consistency.pcm import pcm_inference_sigmas + + return pcm_inference_sigmas( + num_inference_steps, + int(pcm_solver_steps), + scheduler=self.scheduler, + latent_hw=latent_hw, + ) + def _load_infer_sample(self, index, prompt): infer_sample = self.dataloader_eval.dataset[index] infer_sample["conditioning"]["prompt"] = prompt @@ -40,7 +45,12 @@ def _load_dummy_sample(self, records): @torch.no_grad() def infer(self): - records = self.dataloader_eval.dataset.samples + dataset = self.dataloader_eval.dataset + sample_processor = getattr(dataset, "sample_processor", None) + if sample_processor is None: + raise ValueError("image_infer requires a dataset with a sample_processor") + + records = dataset.samples prompts = [record["prompt"] for record in records] rank = get_rank() world_size = get_world_size() @@ -65,7 +75,7 @@ def infer(self): self.guidance_scale = self.infer_config.get("cfg_guidance_scale", 4.0) negative_prompt = self.infer_config.get("negative_prompt", " ") negative_sample = {"inputs": {}, "conditioning": {"prompt": negative_prompt}, "meta": {}} - static_neg_cond = None if has_source_condition else self.model.encode_condition(negative_sample) + static_neg_cond = None if has_source_condition else self.model.encode_inference_condition(negative_sample, is_negative=True) else: self.guidance_scale = None negative_prompt = None @@ -82,10 +92,10 @@ def infer(self): prompt = prompts[i] if has_sample else " " infer_sample = self._load_infer_sample(i, prompt) if has_sample else self._load_dummy_sample(records) - height, width = _target_hw_for_sample(infer_sample, default_height, default_width) + height, width = sample_processor.infer_target_size(infer_sample, default_height, default_width) seed = base_seed + i if has_sample else base_seed generator = torch.Generator(device=self.model.device).manual_seed(seed) - pos_cond = self.model.encode_condition(infer_sample) + pos_cond = self.model.encode_inference_condition(infer_sample) if self.enable_cfg: if has_source_condition: neg_sample = { @@ -93,14 +103,21 @@ def infer(self): "conditioning": {**infer_sample["conditioning"], "prompt": negative_prompt}, "meta": infer_sample["meta"], } - neg_cond = self.model.encode_condition(neg_sample) + neg_cond = self.model.encode_inference_condition(neg_sample, is_negative=True) else: neg_cond = static_neg_cond else: neg_cond = None latent = self.model.prepare_infer_latents(height, width, generator) latent_hw = (latent.shape[-2], latent.shape[-1]) - self.scheduler.set_timesteps(num_inference_steps, latent_hw=latent_hw) + self.scheduler.set_timesteps( + num_inference_steps, + sigmas=self._inference_sigmas( + num_inference_steps, + latent_hw=latent_hw, + ), + latent_hw=latent_hw, + ) total_steps = len(self.scheduler.infer_timesteps) if has_sample: diff --git a/lightx2v_train/lightx2v_train/infer/image_native.py b/lightx2v_train/lightx2v_train/infer/image_native.py deleted file mode 100644 index ad796921c..000000000 --- a/lightx2v_train/lightx2v_train/infer/image_native.py +++ /dev/null @@ -1,58 +0,0 @@ -from pathlib import Path - -import torch -from loguru import logger - -from lightx2v_train.runtime.distributed import is_distributed -from lightx2v_train.utils.registry import INFERENCER_REGISTER - -from .base import BaseInferencer - - -@INFERENCER_REGISTER("image_native_infer") -class NativeImageInferencer(BaseInferencer): - @torch.no_grad() - def infer(self): - if is_distributed(): - raise NotImplementedError("image_native_infer is not supported with torchrun. Use image_infer instead.") - - prompts = [record["prompt"] for record in self.dataloader_eval.dataset.samples] - enable_cfg = self.infer_config.get("enable_cfg", False) - negative_prompt = self.infer_config.get("negative_prompt", " ") if enable_cfg else None - base_seed = self.infer_config.get("seed", 42) - - # Model-specific kwargs (e.g. QwenImage uses `true_cfg_scale` instead of `guidance_scale`) - pipeline_kwargs = self.model.get_pipeline_infer_kwargs(self.infer_config) - - # Use the pipeline's original pretrained scheduler for bit-exact alignment with diffusers - pipe = self.model.assemble_pipeline() - - lora_config = self.infer_config.get("lora_config", None) - lora_path = lora_config.get("path", None) if lora_config else None - if lora_path: - pipe.load_lora_weights(lora_path) - - saved_paths = [] - self.model.set_denoiser_eval() - with torch.no_grad(): - for i, prompt in enumerate(prompts): - generator = torch.Generator(device=self.model.device).manual_seed(base_seed + i) - sample_kwargs = self.model.get_pipeline_sample_kwargs(self.dataloader_eval.dataset[i]) - result = pipe( - prompt=prompt, - negative_prompt=negative_prompt, - generator=generator, - **sample_kwargs, - **pipeline_kwargs, - ) - - if self.output_infer_dir is not None: - save_path = Path(self.output_infer_dir) / f"{i:05d}.png" - result.images[0].save(save_path) - logger.info("Saved to {}", save_path) - saved_paths.append(str(save_path)) - - if lora_path: - pipe.unload_lora_weights() - - return saved_paths diff --git a/lightx2v_train/lightx2v_train/infer/video.py b/lightx2v_train/lightx2v_train/infer/video.py index 976018d72..a11db82fa 100644 --- a/lightx2v_train/lightx2v_train/infer/video.py +++ b/lightx2v_train/lightx2v_train/infer/video.py @@ -32,7 +32,6 @@ def _target_hw_for_sample(sample, default_height, default_width): return default_height, default_width -@INFERENCER_REGISTER("wan_ti2v_5b_infer") @INFERENCER_REGISTER("wan_t2v_14b_infer") @INFERENCER_REGISTER("wan_t2v_infer") class WanT2VInferencer(BaseInferencer): @@ -425,7 +424,6 @@ def infer(self): return saved_paths -@INFERENCER_REGISTER("wan_ti2v_5b_ar_infer") @INFERENCER_REGISTER("wan_t2v_14b_ar_infer") @INFERENCER_REGISTER("wan_t2v_ar_infer") class WanT2VARInferencer(BaseInferencer): @@ -802,4 +800,4 @@ def _frame_seq_length(self, latent): return latent_height * latent_width // (patch_h * patch_w) def _num_frame_per_chunk(self): - return int(getattr(self.model, "num_frame_per_chunk", self.config.get("training", {}).get("teacher_forcing", {}).get("num_frame_per_chunk", 1))) + return int(getattr(self.model, "num_frame_per_chunk", 1)) diff --git a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py new file mode 100644 index 000000000..d0c13ec85 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -0,0 +1,45 @@ +from .base import ( + BoundCapability, + CapabilityNotSupportedError, + CapabilityProvider, + CapabilityRegistry, + ModelCapability, +) +from .training import ( + AutoregressiveDistributionMatchingCapability, + AutoregressiveRolloutContext, + CheckpointCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, + DistributionMatchingProfile, + FlowMatchingSFTCapability, + LossResult, + ParallelCapability, + SFTStepContext, + TeacherForcingCapability, + TeacherForcingStepContext, + TrainableModelCapability, + TrainingCachePayload, +) + +__all__ = [ + "AutoregressiveDistributionMatchingCapability", + "AutoregressiveRolloutContext", + "BoundCapability", + "CapabilityNotSupportedError", + "CapabilityProvider", + "CapabilityRegistry", + "CheckpointCapability", + "ConsistencyModelCapability", + "DistributionMatchingCapability", + "DistributionMatchingProfile", + "FlowMatchingSFTCapability", + "LossResult", + "ModelCapability", + "ParallelCapability", + "SFTStepContext", + "TeacherForcingCapability", + "TeacherForcingStepContext", + "TrainingCachePayload", + "TrainableModelCapability", +] diff --git a/lightx2v_train/lightx2v_train/model_capabilities/base.py b/lightx2v_train/lightx2v_train/model_capabilities/base.py new file mode 100644 index 000000000..b13dc8795 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/base.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import weakref +from abc import ABC +from typing import TypeVar, cast + + +class ModelCapability(ABC): + """Root interface for behavior exposed by a loaded model.""" + + +CapabilityT = TypeVar("CapabilityT", bound=ModelCapability) + + +class CapabilityNotSupportedError(RuntimeError): + """Raised when an algorithm is paired with an incompatible model.""" + + +class CapabilityRegistry: + def __init__(self) -> None: + self._capabilities: dict[type[ModelCapability], ModelCapability] = {} + + def register( + self, + interface: type[CapabilityT], + implementation: CapabilityT, + ) -> None: + if not isinstance(implementation, interface): + raise TypeError(f"{type(implementation).__name__} does not implement {interface.__name__}.") + if interface in self._capabilities: + raise RuntimeError(f"Capability already registered: {interface.__name__}.") + self._capabilities[interface] = implementation + + def supports(self, interface: type[ModelCapability]) -> bool: + return interface in self._capabilities + + def require(self, interface: type[CapabilityT]) -> CapabilityT: + implementation = self._capabilities.get(interface) + if implementation is None: + raise CapabilityNotSupportedError(f"Missing capability: {interface.__name__}.") + return cast(CapabilityT, implementation) + + def missing( + self, + interfaces: tuple[type[ModelCapability], ...], + ) -> tuple[type[ModelCapability], ...]: + return tuple(interface for interface in interfaces if not self.supports(interface)) + + +class CapabilityProvider: + """Base for model wrappers that explicitly publish loaded capabilities.""" + + def __init__(self) -> None: + self.capabilities = CapabilityRegistry() + self._capabilities_registered = False + + def ensure_capabilities(self) -> CapabilityRegistry: + if not self._capabilities_registered: + self.capabilities = CapabilityRegistry() + try: + self.register_capabilities() + except Exception: + self.capabilities = CapabilityRegistry() + raise + self._capabilities_registered = True + return self.capabilities + + def register_capabilities(self) -> None: + """Register capabilities supported by the loaded model instance.""" + + +class BoundCapability(ModelCapability): + """Capability adapter that references, but never owns or copies, a model.""" + + def __init__(self, model: CapabilityProvider) -> None: + self._model_ref = weakref.ref(model) + + @property + def model(self): + model = self._model_ref() + if model is None: + raise RuntimeError("Capability owner has been released.") + return model diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py new file mode 100644 index 000000000..d12f8ff86 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +from abc import abstractmethod +from collections.abc import Collection +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, TypedDict + +import torch +from torch import Tensor + +from .base import ModelCapability + + +@dataclass(frozen=True) +class LossResult: + loss: Tensor + metrics: Mapping[str, Tensor | float] = field(default_factory=dict) + + +class TrainingCachePayload(TypedDict): + inputs: Mapping[str, Any] + conditioning: Mapping[str, Any] + meta: Mapping[str, Any] + + +@dataclass(frozen=True) +class SFTStepContext: + noise_scheduler: Any + running_dtype: torch.dtype + broadcast: Callable[[Any], Any] + + +class FlowMatchingSFTCapability(ModelCapability): + def encode_training_cache(self, batch: Mapping[str, Any]) -> TrainingCachePayload: + """Encode the static sample data consumed by flow-matching training.""" + raise NotImplementedError(f"{type(self).__name__} does not support training-cache encoding.") + + @abstractmethod + def compute_loss( + self, + batch: Mapping[str, Any], + context: SFTStepContext, + ) -> LossResult: + """Compute a differentiable supervised flow-matching loss.""" + + +@dataclass(frozen=True) +class TeacherForcingStepContext: + scheduler: Any + running_dtype: torch.dtype + num_frame_per_chunk: int + noise_augmentation_max_timestep: int + broadcast: Callable[[Any], Any] + + +class TeacherForcingCapability(ModelCapability): + def encode_training_cache(self, batch: Mapping[str, Any]) -> TrainingCachePayload: + """Encode the static target and condition used by teacher forcing.""" + raise NotImplementedError(f"{type(self).__name__} does not support training-cache encoding.") + + @abstractmethod + def compute_loss( + self, + batch: Mapping[str, Any], + context: TeacherForcingStepContext, + ) -> LossResult: + """Compute a differentiable teacher-forcing loss.""" + + +class ConsistencyModelCapability(ModelCapability): + def encode_training_cache(self, batch: Mapping[str, Any]) -> TrainingCachePayload: + """Encode the static sample data consumed by consistency training.""" + raise NotImplementedError(f"{type(self).__name__} does not support training-cache encoding.") + + @abstractmethod + def configure(self, features: Collection[str]) -> None: + """Install model structure required by the selected objective.""" + + @abstractmethod + def restore_trainable_auxiliary(self) -> None: + """Make objective-specific parameters trainable after backbone setup.""" + + @abstractmethod + def auxiliary_parameter_names(self) -> tuple[str, ...]: + """Return configured denoiser keys stored beside parameter-efficient weights.""" + + @abstractmethod + def encode_latent(self, batch): + """Encode one training sample into the objective's clean state.""" + + @abstractmethod + def encode_condition(self, batch): + """Encode one training sample's conditioning inputs.""" + + @abstractmethod + def sampling_latent_hw(self, batch, clean) -> tuple[int, int]: + """Return latent spatial geometry used by resolution-aware schedulers.""" + + @abstractmethod + def predict(self, request, path): + """Run the model and convert its output to the requested parameterization.""" + + @abstractmethod + def predict_log_variance(self, time): + """Predict the learned loss log-variance required by sCM.""" + + @abstractmethod + def set_frozen(self, training: bool = False) -> None: + """Freeze the denoiser while selecting its train/eval execution mode.""" + + @abstractmethod + def denoiser(self): + """Return the module owned by this capability.""" + + +@dataclass(frozen=True) +class DistributionMatchingProfile: + """Feature and precision constraints published by a distribution-matching adapter.""" + + supports_guidance: bool = True + supports_ida: bool = True + supports_diversity: bool = True + supports_real_data_fake: bool = True + supports_warped_denoising_schedule: bool = True + default_latent_dtype: torch.dtype | None = None + # ``None`` keeps the capability available to every DMD-family trainer. + supported_training_methods: frozenset[str] | None = None + + +class DistributionMatchingCapability(ModelCapability): + """Model operations required by distribution-matching algorithms.""" + + @property + def profile(self) -> DistributionMatchingProfile: + return DistributionMatchingProfile() + + def encode_training_cache(self, batch: Mapping[str, Any]) -> TrainingCachePayload: + """Encode the static sample data consumed by distribution matching.""" + raise NotImplementedError(f"{type(self).__name__} does not support training-cache encoding.") + + @property + @abstractmethod + def device(self) -> torch.device: + pass + + @property + @abstractmethod + def default_negative_prompt(self) -> str | None: + pass + + @property + @abstractmethod + def default_lora_target_modules(self): + pass + + @abstractmethod + def latent_shape( + self, + batch: Mapping[str, Any], + shape_config: Mapping[str, Any], + image_sizes, + broadcast: Callable[[Any], Any], + ): + pass + + @abstractmethod + def encode_conditions( + self, + batch: Mapping[str, Any], + negative_prompt: str | None, + guidance_scale: float, + broadcast: Callable[[Any], Any], + ): + pass + + @abstractmethod + def predict_velocity(self, latents, sigma, condition): + pass + + @abstractmethod + def predict_guided_velocity( + self, + latents, + sigma, + condition, + negative_condition, + guidance_scale, + cfg_norm, + ): + pass + + @abstractmethod + def initial_latents(self, latent_shape, dtype, broadcast): + pass + + @abstractmethod + def latent_hw(self, latent_shape): + pass + + @abstractmethod + def random_noise_like(self, latents, dtype, broadcast): + pass + + @abstractmethod + def add_noise(self, scheduler, latents, noise, sigma): + pass + + @abstractmethod + def training_target(self, latents, noise): + pass + + @abstractmethod + def step(self, scheduler, velocity, step_index, sample): + pass + + @abstractmethod + def x0_from_velocity(self, sample, velocity, sigma): + pass + + @abstractmethod + def regression_loss(self, prediction, target): + pass + + @abstractmethod + def dmd_loss(self, latents, fake_x0, teacher_x0): + pass + + @abstractmethod + def detach(self, value): + pass + + @abstractmethod + def to_dtype(self, value, dtype): + pass + + @abstractmethod + def extract_real_latents(self, batch, dtype, broadcast): + pass + + @abstractmethod + def set_training(self, enabled: bool) -> None: + pass + + @abstractmethod + def denoiser(self): + pass + + +@dataclass(frozen=True) +class AutoregressiveRolloutContext: + denoising_steps: Tensor + denoising_scheduler: Any + trajectory_scheduler: Any + running_dtype: torch.dtype + frames_per_chunk: int + same_step_across_blocks: bool + context_noise: float + sequence_parallel_cache: bool + grad_enabled: bool + + +class AutoregressiveDistributionMatchingCapability(ModelCapability): + @abstractmethod + def rollout( + self, + condition, + latent_shape, + initial_latents, + context: AutoregressiveRolloutContext, + ): + """Run a differentiable chunk-wise autoregressive rollout.""" + + +class TrainableModelCapability(ModelCapability): + @abstractmethod + def configure(self, train_type: str, lora_config: Mapping[str, Any]) -> None: + pass + + @abstractmethod + def restore(self, train_type: str) -> None: + pass + + @abstractmethod + def parameters(self): + pass + + @abstractmethod + def enable_gradient_checkpointing(self) -> None: + pass + + @abstractmethod + def set_eval(self) -> None: + pass + + @abstractmethod + def log_structure(self) -> None: + pass + + +class ParallelCapability(ModelCapability): + @abstractmethod + def apply(self, config) -> None: + pass + + @abstractmethod + def set_gradient_sync(self, enabled: bool) -> None: + pass + + @abstractmethod + def is_fsdp(self) -> bool: + pass + + @abstractmethod + def state_module(self): + pass + + +class CheckpointCapability(ModelCapability): + @abstractmethod + def save_weights(self, save_dir, train_type) -> None: + pass + + @abstractmethod + def load_weights(self, save_dir, train_type) -> None: + pass + + @abstractmethod + def save_consolidated(self, output_path) -> None: + pass + + @abstractmethod + def save_full_model(self, output_path) -> None: + pass diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 268dd9e8a..737173db5 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -3,20 +3,38 @@ from lightx2v_train.utils.registry import build_model _LAZY_EXPORTS = { - "Flux2DevModel": (".flux2_dev", "Flux2DevModel"), - "Flux2KleinModel": (".flux2_klein", "Flux2KleinModel"), - "LingBotVideoModel": (".lingbot_video", "LingBotVideoModel"), - "LongCatImageModel": (".longcat_image", "LongCatImageModel"), - "LTX2T2AVModel": (".ltx_t2av", "LTX2T2AVModel"), - "MiniMaxH3T2AVModel": (".minimax_h3_t2av", "MiniMaxH3T2AVModel"), - "QwenImageModel": (".qwen_image", "QwenImageModel"), - "QwenImageEditModel": (".qwen_image_edit", "QwenImageEditModel"), - "WanFastWAMModel": (".wan_fastwam", "WanFastWAMModel"), - "WanT2VModel": (".wan_t2v", "WanT2VModel"), - "WanTI2V5BModel": (".wan_ti2v_5b", "WanTI2V5BModel"), + "Flux2DevModel": (".flux2.flux2_dev", "Flux2DevModel"), + "Flux2DevEditModel": (".flux2.flux2_dev_edit", "Flux2DevEditModel"), + "Flux2KleinModel": (".flux2.flux2_klein", "Flux2KleinModel"), + "Flux2KleinEditModel": (".flux2.flux2_klein_edit", "Flux2KleinEditModel"), + "LingBotVideoModel": (".wan.lingbot_video", "LingBotVideoModel"), + "LongCatImageModel": (".longcat_image.longcat_image", "LongCatImageModel"), + "LongCatImageEditModel": (".longcat_image.longcat_image_edit", "LongCatImageEditModel"), + "MiniMaxH3T2AVModel": (".minimax_h3.minimax_h3_t2av", "MiniMaxH3T2AVModel"), + "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), + "QwenImageEditModel": (".qwen_image.qwen_image_edit", "QwenImageEditModel"), + "WanT2VModel": (".wan.wan_t2v", "WanT2VModel"), } +def build_loaded_model( + config, + *, + load_transformer, + load_vae, + load_condition_encoder, +): + """Build a model wrapper, load its components, then publish capabilities.""" + model = build_model(config) + model.load_components( + load_transformer=load_transformer, + load_vae=load_vae, + load_condition_encoder=load_condition_encoder, + ) + model.ensure_capabilities() + return model + + def __getattr__(name): target = _LAZY_EXPORTS.get(name) if target is None: @@ -29,15 +47,16 @@ def __getattr__(name): __all__ = [ "build_model", + "build_loaded_model", "QwenImageModel", "QwenImageEditModel", "LongCatImageModel", + "LongCatImageEditModel", "Flux2DevModel", + "Flux2DevEditModel", "Flux2KleinModel", + "Flux2KleinEditModel", "LingBotVideoModel", - "LTX2T2AVModel", "MiniMaxH3T2AVModel", "WanT2VModel", - "WanTI2V5BModel", - "WanFastWAMModel", ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 0419e9eaa..9c3d08e68 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -12,23 +12,66 @@ from safetensors.torch import load_file, save_file from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict +from lightx2v_train.model_capabilities import ( + CapabilityProvider, + CheckpointCapability, + ParallelCapability, + TrainableModelCapability, +) +from lightx2v_train.model_zoo.capability_adapters.common import ( + CommonCheckpointCapability, + CommonParallelCapability, + CommonTrainableCapability, +) from lightx2v_train.runtime.distributed import is_main_process from lightx2v_train.runtime.fsdp import is_fsdp2_module from lightx2v_train.utils.utils import get_running_dtype -class BaseModel: +class BaseModel(CapabilityProvider): + default_unconditional_prompt = " " + shared_condition_keys = () + def __init__(self, config): + super().__init__() self.config = config self.running_dtype = get_running_dtype(config["model"]["running_dtype"]) self.device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") self.vae = None + self.vae_config = None + self.text_pipeline = None + self.image_processor = None + + def register_capabilities(self): + self.capabilities.register( + TrainableModelCapability, + CommonTrainableCapability(self), + ) + self.capabilities.register( + ParallelCapability, + CommonParallelCapability(self), + ) + self.capabilities.register( + CheckpointCapability, + CommonCheckpointCapability(self), + ) - def load_components(self, transformer_only=False, reference_model=None): + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + """Load the requested model weights; lightweight config metadata may still be read.""" raise NotImplementedError - def dmd_latent_shape(self, batch_size, height, width): - raise NotImplementedError(f"{self.__class__.__name__} must define dmd_latent_shape().") + def reuse_frozen_components_from(self, source): + """Reuse the frozen VAE and condition components owned by another model.""" + self.vae = source.vae + self.vae_config = source.vae_config + self.text_pipeline = source.text_pipeline + self.image_processor = source.image_processor def denoiser_module(self): raise NotImplementedError(f"{self.__class__.__name__} must define denoiser_module().") @@ -42,64 +85,6 @@ def add_lora(self, rank, alpha, target_modules): ) self.denoiser_module().add_adapter(lora_config) - def add_dual_lora( - self, - rank, - alpha, - target_modules, - student_adapter="student", - teacher_adapter="teacher", - init_teacher_from_student=True, - ): - lora_config = LoraConfig( - r=rank, - lora_alpha=alpha, - init_lora_weights="gaussian", - target_modules=target_modules, - ) - denoiser = self.denoiser_module() - denoiser.requires_grad_(False) - denoiser.add_adapter(lora_config, adapter_name=student_adapter) - denoiser.add_adapter(lora_config, adapter_name=teacher_adapter) - denoiser.set_adapter(student_adapter) - if init_teacher_from_student: - self.copy_lora_adapter_weights(student_adapter, teacher_adapter) - - @torch.no_grad() - def copy_lora_adapter_weights(self, src_adapter, dst_adapter): - named_params = dict(self.denoiser_module().named_parameters()) - for name, param in named_params.items(): - if src_adapter not in name: - continue - dst_name = name.replace(src_adapter, dst_adapter) - if dst_name in named_params: - named_params[dst_name].data.copy_(param.data) - - def set_active_adapter(self, adapter_name): - self.denoiser_module().set_adapter(adapter_name) - - def set_dual_lora_trainable(self, student_adapter="student", teacher_adapter="teacher"): - denoiser = self.denoiser_module() - denoiser.requires_grad_(False) - denoiser.train() - for name, param in denoiser.named_parameters(): - if student_adapter in name and "lora" in name: - param.requires_grad = True - else: - param.requires_grad = False - - @torch.no_grad() - def ema_update_lora_adapter(self, src_adapter="student", dst_adapter="teacher", ema_decay=0.999): - named_params = dict(self.denoiser_module().named_parameters()) - for name, src_param in named_params.items(): - if src_adapter not in name: - continue - dst_name = name.replace(src_adapter, dst_adapter) - if dst_name not in named_params: - continue - dst_param = named_params[dst_name] - dst_param.data.mul_(ema_decay).add_(src_param.data, alpha=1.0 - ema_decay) - def set_lora_trainable(self): denoiser = self.denoiser_module() denoiser.requires_grad_(False) @@ -154,6 +139,40 @@ def encode_to_latent(self, sample): def encode_condition(self, sample): raise NotImplementedError + @property + def unconditional_prompt(self): + return self.config["model"].get("unconditional_prompt", self.default_unconditional_prompt) + + def encode_condition_roles(self, sample, prompts, *, contextual_roles=()): + """Encode named prompts, retaining sample context only for requested roles.""" + contextual_roles = set(contextual_roles) + contextual_names = [name for name in prompts if name in contextual_roles] + conditions = {name: self.encode_prompt_condition(prompt) for name, prompt in prompts.items() if name not in contextual_roles} + if contextual_names: + contextual_prompts = [prompts[name] for name in contextual_names] + contextual_conditions = self.encode_conditions_with_context(sample, contextual_prompts) + conditions.update(zip(contextual_names, contextual_conditions, strict=True)) + return {name: conditions[name] for name in prompts} + + def encode_conditions_with_context(self, sample, prompts): + conditions = [] + for prompt in prompts: + contextual_sample = dict(sample) + contextual_sample["conditioning"] = { + **sample["conditioning"], + "prompt": prompt, + } + conditions.append(self.encode_condition(contextual_sample)) + return conditions + + def encode_to_cache_latent(self, sample): + """Encode a deterministic target latent for reuse across training epochs.""" + return self.encode_to_latent(sample) + + def encode_inference_condition(self, sample, *, is_negative=False): + del is_negative + return self.encode_condition(sample) + def prepare_denoiser_input(self, noisy_latent, condition=None): raise NotImplementedError @@ -163,34 +182,43 @@ def denoise(self, denoiser_input, timesteps, condition): def postprocess_denoiser_output(self, prediction, denoiser_input): raise NotImplementedError + def apply_cfg(self, positive, negative, guidance_scale): + return negative + guidance_scale * (positive - negative) + + def denoiser_prediction_type(self): + """Return the quantity predicted by the denoiser. + + LightX2V's current diffusion backbones are trained as rectified-flow + velocity predictors. Keeping this declaration on the model avoids + baking that assumption into training objectives and leaves room + for models that predict x0, noise, or another parameterization. + """ + return "velocity" + + def predict_denoiser_output(self, noisy_latent, timestep_or_sigma, condition, **denoiser_kwargs): + """Run the model-specific denoiser path and return latent-shaped output. + + Consistency objectives operate on latent tensors, while individual + models may pack those tensors before the transformer forward. This + method is the common boundary between the two layers. Extra keyword + arguments are intentionally forwarded for algorithms such as + MeanFlow, whose denoisers can require an additional endpoint time. + """ + denoiser_input = self.prepare_denoiser_input(noisy_latent, condition=condition) + prediction = self.denoise( + denoiser_input, + timestep_or_sigma, + condition, + **denoiser_kwargs, + ) + return self.postprocess_denoiser_output(prediction, denoiser_input) + def prepare_infer_latents(self, height, width, generator=None): raise NotImplementedError - def dmd_latent_shape(self, batch_size, height, width): - raise NotImplementedError(f"{self.__class__.__name__} must define dmd_latent_shape().") - - def cfg_on_denoiser_output(self): - return False - def decode_latent(self, latent): raise NotImplementedError - def assemble_pipeline(self, scheduler=None): - raise NotImplementedError - - def get_pipeline_infer_kwargs(self, infer_config): - """Return kwargs to pass to pipeline.__call__. Override to adapt model-specific parameter names.""" - return { - "height": infer_config.get("height", 1024), - "width": infer_config.get("width", 1024), - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", 4.0), - } - - def get_pipeline_sample_kwargs(self, sample): - """Return per-sample kwargs to pass to pipeline.__call__ during native inference.""" - return {} - def load_lora_for_infer(self, lora_path, adapter_name=None): denoiser = self.denoiser_module() if adapter_name is None: @@ -204,8 +232,19 @@ def unload_lora_for_infer(self): self.denoiser_module().delete_adapters(adapter_name) self._infer_lora_adapter_name = None - def save_lora_weights(self, save_dir, adapter_name=None, weights_subdir=None): - peft_state_dict = self._get_lora_state_dict_for_save(adapter_name=adapter_name) + def save_lora_weights( + self, + save_dir, + adapter_name=None, + weights_subdir=None, + *, + auxiliary_parameter_names=(), + auxiliary_weights_name=None, + ): + peft_state_dict, auxiliary_state_dict = self._get_lora_and_auxiliary_state_dict_for_save( + adapter_name=adapter_name, + auxiliary_parameter_names=auxiliary_parameter_names, + ) if not is_main_process(): return @@ -216,23 +255,44 @@ def save_lora_weights(self, save_dir, adapter_name=None, weights_subdir=None): self.pipeline_cls.save_lora_weights(output_dir, lora_state_dict, safe_serialization=True) else: save_file(lora_state_dict, os.path.join(output_dir, "pytorch_lora_weights.safetensors")) + if auxiliary_state_dict: + if not auxiliary_weights_name: + raise ValueError("auxiliary_weights_name is required when auxiliary parameters are saved.") + save_file( + auxiliary_state_dict, + os.path.join(output_dir, auxiliary_weights_name), + ) def _get_lora_state_dict_for_save(self, adapter_name=None): + return self._get_lora_and_auxiliary_state_dict_for_save(adapter_name=adapter_name)[0] + + def _get_lora_and_auxiliary_state_dict_for_save( + self, + adapter_name=None, + auxiliary_parameter_names=(), + ): denoiser = self.denoiser_module() peft_kwargs = {} if adapter_name is None else {"adapter_name": adapter_name} if not is_fsdp2_module(denoiser): - return get_peft_model_state_dict(denoiser, **peft_kwargs) + state_dict = denoiser.state_dict() + else: + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + ignore_frozen_params=False, + strict=False, + ) + state_dict, _ = get_state_dict(denoiser, (), options=options) + if not is_main_process(): + return {}, {} - options = StateDictOptions( - full_state_dict=True, - cpu_offload=True, - ignore_frozen_params=False, - strict=False, - ) - state_dict, _ = get_state_dict(denoiser, (), options=options) - if not is_main_process(): - return {} - return get_peft_model_state_dict(denoiser, state_dict=state_dict, **peft_kwargs) + peft_state_dict = get_peft_model_state_dict(denoiser, state_dict=state_dict, **peft_kwargs) + auxiliary_names = set(auxiliary_parameter_names) + missing = auxiliary_names - state_dict.keys() + if missing: + raise RuntimeError(f"Auxiliary parameters are missing from the model state: {sorted(missing)}") + auxiliary_state_dict = {name: state_dict[name].detach().cpu().contiguous() for name in auxiliary_names} + return peft_state_dict, auxiliary_state_dict def load_lora_weights_for_resume(self, lora_path, adapter_name=None, weights_subdir=None): weights_dir = os.path.join(lora_path, weights_subdir) if weights_subdir else lora_path @@ -249,6 +309,27 @@ def load_lora_weights_for_resume(self, lora_path, adapter_name=None, weights_sub if incompatible and incompatible.unexpected_keys: logger.warning("Unexpected keys when resuming LoRA: {}", incompatible.unexpected_keys) + def load_auxiliary_weights( + self, + checkpoint_dir, + parameter_names, + *, + weights_name, + ): + names = set(parameter_names) + if not names: + return + path = os.path.join(checkpoint_dir, weights_name) + if not os.path.exists(path): + raise RuntimeError(f"Auxiliary weights were not found at {path}.") + incompatible = self.denoiser_module().load_state_dict(load_file(path), strict=False) + missing = [name for name in incompatible.missing_keys if name in names] + unexpected = [name for name in incompatible.unexpected_keys if name in names] + if missing: + raise RuntimeError(f"Missing auxiliary keys in {path}: {missing}") + if unexpected: + logger.warning("Unexpected auxiliary keys in {}: {}", path, unexpected) + def load_full_weights_for_resume(self, resume_ckpt_path): raise NotImplementedError(f"{self.__class__.__name__} must define load_full_weights_for_resume().") diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py new file mode 100644 index 000000000..110a577b0 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py @@ -0,0 +1,8 @@ +"""Capability adapters shared by multiple model families.""" + +from .latent_geometry import LatentGeometry, SpatialLatentGeometry + +__all__ = [ + "LatentGeometry", + "SpatialLatentGeometry", +] diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py new file mode 100644 index 000000000..04305bb8c --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -0,0 +1,737 @@ +from __future__ import annotations + +import os +from collections.abc import Collection, Mapping +from typing import Any + +import torch +import torch.nn.functional as F + +from lightx2v_train.data.training_cache import preserve_cache_dtype +from lightx2v_train.model_capabilities import ( + BoundCapability, + CheckpointCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, + FlowMatchingSFTCapability, + LossResult, + ParallelCapability, + SFTStepContext, + TrainableModelCapability, +) +from lightx2v_train.runtime.distributed import is_main_process +from lightx2v_train.runtime.parallel import ( + apply_parallel, + set_parallel_gradient_sync, +) +from lightx2v_train.utils.image_size_buckets import parse_image_size_buckets + +from .latent_geometry import LatentGeometry + + +def _require_single_prompt(prompt): + if isinstance(prompt, str): + return True + prompts = list(prompt) + if len(prompts) != 1: + raise ValueError(f"Training requires exactly one prompt per rank; physical batch sizes greater than 1 are not supported, got {len(prompts)} prompts.") + return False + + +def _negative_prompt(conditioning, fallback, scalar=False): + value = conditioning.get("negative_prompt") + if value is None: + prompts = [] + elif isinstance(value, str): + prompts = [value] + else: + prompts = list(value) + if not prompts: + prompt = fallback + elif len(prompts) == 1: + prompt = prompts[0] + else: + raise ValueError(f"Training requires exactly one negative prompt per rank; got {len(prompts)} negative prompts.") + prompt = prompt if isinstance(prompt, str) and prompt.strip() else fallback + return prompt if scalar else [prompt] + + +def _require_singleton_tensor(value, name): + if not torch.is_tensor(value) or value.ndim == 0 or value.shape[0] != 1: + shape = tuple(value.shape) if torch.is_tensor(value) else type(value).__name__ + raise ValueError(f"{name} must have leading dimension 1; physical batch sizes greater than 1 are not supported, got {shape}.") + return value + + +def _move_cached_value(value, model, key=None): + if torch.is_tensor(value): + kwargs = {"device": model.device} + if value.is_floating_point() and not preserve_cache_dtype(key): + kwargs["dtype"] = model.running_dtype + return value.to(**kwargs) + if isinstance(value, Mapping): + return {name: _move_cached_value(item, model, name) for name, item in value.items()} + if isinstance(value, tuple): + return tuple(_move_cached_value(item, model, key) for item in value) + if isinstance(value, list): + return [_move_cached_value(item, model, key) for item in value] + return value + + +def _cached_latent(batch, model): + latent = batch.get("inputs", {}).get("latents") + if latent is not None: + return _move_cached_value(latent, model) + cache_path = batch.get("meta", {}).get("training_cache_path") + if cache_path is not None: + raise KeyError(f"Training cache {cache_path} has no inputs.latents entry.") + return None + + +def _cached_condition(batch, model, role=None): + conditioning = batch.get("conditioning", {}) + if "positive" not in conditioning: + cache_path = batch.get("meta", {}).get("training_cache_path") + if cache_path is not None: + raise KeyError(f"Training cache {cache_path} has no conditioning.positive entry.") + return None + role = role or conditioning.get("active", "positive") + if role not in conditioning: + cache_path = batch.get("meta", {}).get("training_cache_path", "") + raise KeyError(f"Training cache {cache_path} has no conditioning.{role} entry.") + condition = dict(conditioning.get("shared", {})) + condition.update(conditioning[role]) + return _move_cached_value(condition, model) + + +def _cached_condition_pair(batch, model, *, require_negative): + positive = _cached_condition(batch, model) + if positive is None: + return None + negative = _cached_condition(batch, model, role="negative") if require_negative else None + return positive, negative + + +def _prompt_or_default(value, default): + return value if isinstance(value, str) else default + + +def _uses_prompt_dropout(model): + return float(model.config.get("data", {}).get("train", {}).get("prompt_dropout_rate", 0.0)) > 0 + + +def _consistency_requires_negative(model): + config = model.config["training"].get("consistency", {}) + algorithm = str(config.get("algorithm", "cm")).lower() + teacher = config.get("teacher", {}) + guidance_scale = teacher.get("guidance_scale") + if algorithm in {"mean_flow", "meanflow"}: + guidance = config.get("guidance", {}) + return any( + value is not None + for value in ( + guidance.get("condition_dropout_probability"), + guidance.get("scale", guidance_scale), + guidance.get("mixture_ratio"), + ) + ) + if algorithm == "pcm": + return guidance_scale is not None and float(guidance_scale) != 1.0 + return str(config.get("mode", "ct")).lower() == "cd" and guidance_scale is not None and float(guidance_scale) != 1.0 + + +def _target_cache_inputs(model, batch, *, required): + target = batch.get("inputs", {}).get("target_pixel_values") + if target is None: + if required: + raise ValueError("Training-cache encoding requires inputs.target_pixel_values.") + return {} + return {"latents": model.encode_to_cache_latent(batch)} + + +def _encode_cache_conditions(model, batch, prompts, contextual_roles): + conditions = model.encode_condition_roles( + batch, + prompts, + contextual_roles=contextual_roles, + ) + result = {} + for role, condition in conditions.items(): + if not isinstance(condition, Mapping): + raise TypeError(f"Encoded {role} condition must be a mapping, got {type(condition).__name__}.") + result[role] = dict(condition) + return result + + +def _cache_values_equal(left, right): + if left is right: + return True + if torch.is_tensor(left) or torch.is_tensor(right): + return torch.is_tensor(left) and torch.is_tensor(right) and left.shape == right.shape and left.dtype == right.dtype and torch.equal(left, right) + if isinstance(left, Mapping) or isinstance(right, Mapping): + return isinstance(left, Mapping) and isinstance(right, Mapping) and left.keys() == right.keys() and all(_cache_values_equal(left[key], right[key]) for key in left) + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + return type(left) is type(right) and len(left) == len(right) and all(_cache_values_equal(a, b) for a, b in zip(left, right, strict=True)) + return left == right + + +def _factor_shared_conditions(conditions, shared_keys): + shared = {} + for key in shared_keys: + if not all(key in condition for condition in conditions.values()): + continue + values = [condition[key] for condition in conditions.values()] + if not all(_cache_values_equal(values[0], value) for value in values[1:]): + raise ValueError(f"Condition field {key!r} is declared shared but differs between roles.") + shared[key] = values[0] + for condition in conditions.values(): + condition.pop(key) + return shared + + +def _training_cache_data(model, batch, *, inputs, prompts, contextual_roles=(), conditioning_meta=None): + conditions = _encode_cache_conditions(model, batch, prompts, contextual_roles) + shared = _factor_shared_conditions(conditions, model.shared_condition_keys) + conditioning = { + "prompt": batch["conditioning"]["prompt"], + **conditions, + **(conditioning_meta or {}), + } + if "unconditional" in conditions: + conditioning["unconditional_prompt"] = model.unconditional_prompt + if shared: + conditioning["shared"] = shared + + meta = dict(batch.get("meta", {})) + target = batch.get("inputs", {}).get("target_pixel_values") + if target is not None: + meta["target_height"] = int(target.shape[-2]) + meta["target_width"] = int(target.shape[-1]) + return { + "inputs": inputs, + "conditioning": conditioning, + "meta": meta, + } + + +class CommonTrainableCapability(BoundCapability, TrainableModelCapability): + def configure(self, train_type: str, lora_config: Mapping[str, Any]) -> None: + if train_type == "lora": + rank = int(lora_config.get("rank", 16)) + self.model.add_lora( + rank, + int(lora_config.get("alpha", rank)), + lora_config.get("target_modules"), + ) + self.model.set_lora_trainable() + return + if train_type != "full": + raise ValueError(f"Unsupported train type {train_type!r}; expected 'lora' or 'full'.") + self.model.set_full_trainable() + + def restore(self, train_type: str) -> None: + if train_type == "lora": + self.model.set_lora_trainable() + elif train_type == "full": + self.model.set_full_trainable() + else: + raise ValueError(f"Unsupported train type {train_type!r}; expected 'lora' or 'full'.") + + def parameters(self): + return self.model.trainable_parameters() + + def enable_gradient_checkpointing(self) -> None: + self.model.enable_gradient_checkpointing() + + def set_eval(self) -> None: + self.model.set_denoiser_eval() + + def log_structure(self) -> None: + self.model.log_model_structure() + + +class CommonParallelCapability(BoundCapability, ParallelCapability): + def apply(self, config) -> None: + apply_parallel(self.model, config) + + def set_gradient_sync(self, enabled: bool) -> None: + set_parallel_gradient_sync(self.model, enabled) + + def is_fsdp(self) -> bool: + return self.model.is_fsdp2_wrapped() + + def state_module(self): + return self.model.fsdp2_state_module() + + +class CommonCheckpointCapability(BoundCapability, CheckpointCapability): + _CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME = "consistency_auxiliary.safetensors" + + def _consistency_model_auxiliary_parameter_names(self) -> tuple[str, ...]: + capabilities = self.model.ensure_capabilities() + if not capabilities.supports(ConsistencyModelCapability): + return () + return capabilities.require(ConsistencyModelCapability).auxiliary_parameter_names() + + def save_weights(self, save_dir, train_type) -> None: + if train_type == "lora": + self.model.save_lora_weights( + save_dir, + auxiliary_parameter_names=self._consistency_model_auxiliary_parameter_names(), + auxiliary_weights_name=self._CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME, + ) + elif is_main_process(): + torch.save( + self.model.denoiser_module().state_dict(), + os.path.join(save_dir, "model_state.pt"), + ) + + def load_weights(self, save_dir, train_type) -> None: + if train_type == "lora": + self.model.load_lora_weights_for_resume(save_dir) + self.model.load_auxiliary_weights( + save_dir, + self._consistency_model_auxiliary_parameter_names(), + weights_name=self._CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME, + ) + return + path = os.path.join(save_dir, "model_state.pt") + if not os.path.exists(path): + raise RuntimeError(f"model_state.pt not found in {save_dir}") + state_dict = torch.load( + path, + map_location="cpu", + weights_only=False, + ) + self.model.denoiser_module().load_state_dict(state_dict) + + def save_consolidated(self, output_path) -> None: + self.model.save_consolidated_weights(output_path) + + def save_full_model(self, output_path) -> None: + self.model.save_full_model(output_path) + + +class GenericFlowMatchingCapability(BoundCapability, FlowMatchingSFTCapability): + def encode_training_cache(self, batch): + prompt = batch["conditioning"]["prompt"] + prompts = {"positive": prompt} + if _uses_prompt_dropout(self.model): + prompts["unconditional"] = self.model.unconditional_prompt + return _training_cache_data( + self.model, + batch, + inputs=_target_cache_inputs(self.model, batch, required=True), + prompts=prompts, + contextual_roles=prompts, + ) + + def compute_loss( + self, + batch: Mapping[str, Any], + context: SFTStepContext, + ) -> LossResult: + scheduler = context.noise_scheduler + broadcast = context.broadcast + + with torch.no_grad(): + latent = _cached_latent(batch, self.model) + if latent is None: + latent = self.model.encode_to_latent(batch) + latent = _require_singleton_tensor( + broadcast(latent), + "Flow-matching latent", + ) + noise = broadcast(torch.randn_like(latent, dtype=context.running_dtype)) + latent_hw = latent.shape[-2:] + timestep_or_sigma = broadcast( + scheduler.sample_timestep_or_sigma( + latent_hw=latent_hw, + ) + ) + noisy_latent = scheduler.add_noise( + latent, + noise, + timestep_or_sigma, + ) + condition = _cached_condition(batch, self.model) + if condition is None: + condition = self.model.encode_condition(batch) + condition = broadcast(condition) + + prediction = self.model.predict_denoiser_output( + noisy_latent, + timestep_or_sigma, + condition, + ) + target = scheduler.build_train_gt(latent, noise) + loss = (prediction.float() - target.float()).square().mean() + return LossResult(loss=loss) + + +class GenericConsistencyModelCapability(BoundCapability, ConsistencyModelCapability): + def encode_training_cache(self, batch): + prompt = batch["conditioning"]["prompt"] + teacher = self.model.config["training"].get("consistency", {}).get("teacher", {}) + negative_prompt = _prompt_or_default(teacher.get("negative_prompt"), self.model.unconditional_prompt) + prompts = {"positive": prompt} + conditioning_meta = {} + if _consistency_requires_negative(self.model): + prompts["negative"] = negative_prompt + conditioning_meta["negative_prompt"] = negative_prompt + if _uses_prompt_dropout(self.model): + prompts["unconditional"] = self.model.unconditional_prompt + return _training_cache_data( + self.model, + batch, + inputs=_target_cache_inputs(self.model, batch, required=True), + prompts=prompts, + contextual_roles=prompts, + conditioning_meta=conditioning_meta, + ) + + def configure(self, features: Collection[str]) -> None: + features = frozenset(features) + if features: + names = ", ".join(sorted(features)) + raise NotImplementedError(f"{type(self.model).__name__} does not support consistency features: {names}.") + + def restore_trainable_auxiliary(self) -> None: + pass + + def auxiliary_parameter_names(self) -> tuple[str, ...]: + return () + + def encode_latent(self, batch): + latent = _cached_latent(batch, self.model) + if latent is None: + latent = self.model.encode_to_latent(batch) + return _require_singleton_tensor( + latent, + "Consistency latent", + ) + + def encode_condition(self, batch): + condition = _cached_condition(batch, self.model) + return self.model.encode_condition(batch) if condition is None else condition + + def sampling_latent_hw(self, batch, clean) -> tuple[int, int]: + del batch + return int(clean.shape[-2]), int(clean.shape[-1]) + + def predict(self, request, path): + prediction = self.model.predict_denoiser_output( + request.sample, + request.time, + request.condition, + **request.model_kwargs, + ) + return path.convert_prediction( + request.sample, + prediction, + request.time, + source_type=self.model.denoiser_prediction_type(), + target_type=request.prediction_type, + ) + + def predict_log_variance(self, time): + del time + raise NotImplementedError(f"{type(self.model).__name__} does not provide a consistency log-variance head.") + + def set_frozen(self, training: bool = False) -> None: + denoiser = self.model.denoiser_module() + denoiser.requires_grad_(False) + denoiser.train(training) + + def denoiser(self): + return self.model.denoiser_module() + + +class GenericDistributionMatchingCapability(BoundCapability, DistributionMatchingCapability): + """Distribution-matching operations shared by image flow models.""" + + cache_uses_sample_context = False + + def __init__( + self, + model, + *, + latent_geometry: LatentGeometry | None = None, + guidance_in_denoiser_space: bool = False, + ) -> None: + super().__init__(model) + self._latent_geometry = latent_geometry + self._guidance_in_denoiser_space = bool(guidance_in_denoiser_space) + + @property + def device(self): + return self.model.device + + @property + def default_negative_prompt(self): + return self.model.unconditional_prompt + + @property + def default_lora_target_modules(self): + return None + + def encode_training_cache(self, batch): + training = self.model.config["training"] + teacher = training.get("teacher", {}) + dmd = training.get("dmd", {}) + guidance_scale = float(teacher.get("guidance_scale", dmd.get("guidance_scale", 3.0))) + negative_prompt = _prompt_or_default( + teacher.get( + "negative_prompt", + dmd.get("negative_prompt", self.default_negative_prompt), + ), + self.model.unconditional_prompt, + ) + prompt = batch["conditioning"]["prompt"] + prompts = {"positive": prompt} + if _uses_prompt_dropout(self.model): + prompts["unconditional"] = self.model.unconditional_prompt + conditioning_meta = {} + if guidance_scale > 1.0: + prompts["negative"] = negative_prompt + conditioning_meta["negative_prompt"] = negative_prompt + return _training_cache_data( + self.model, + batch, + inputs=_target_cache_inputs(self.model, batch, required=False), + prompts=prompts, + contextual_roles=prompts if self.cache_uses_sample_context else (), + conditioning_meta=conditioning_meta, + ) + + def latent_shape( + self, + batch, + shape_config, + image_sizes, + broadcast, + ): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + _require_single_prompt(prompt) + meta = batch.get("meta", {}) + height = self._target_dimension( + meta, + "target_height", + ) + width = self._target_dimension( + meta, + "target_width", + ) + height = int(broadcast(height)) + width = int(broadcast(width)) + + if image_sizes: + configured_buckets = {bucket.spatial_size for bucket in parse_image_size_buckets(image_sizes)} + if (height, width) not in configured_buckets: + configured = ", ".join(f"{bucket_height}x{bucket_width}" for bucket_height, bucket_width in sorted(configured_buckets)) + raise ValueError(f"Image DMD sample size {height}x{width} is not in training.dmd.image_sizes: [{configured}].") + + if self._latent_geometry is None: + raise NotImplementedError(f"{type(self).__name__} must override latent_shape() or be constructed with a latent_geometry adapter.") + return self._latent_geometry.shape(self.model, height, width) + + @staticmethod + def _target_dimension(meta, key): + if key not in meta: + raise KeyError(f"Image DMD prompt-only samples require meta.target_height and meta.target_width; missing {key}. target_image is not used.") + value = meta[key] + if torch.is_tensor(value): + values = value.detach().reshape(-1).tolist() + elif isinstance(value, (list, tuple)): + values = [item.item() if torch.is_tensor(item) else item for item in value] + else: + values = [value] + values = [int(item) for item in values] + if len(values) != 1: + raise ValueError(f"Image DMD requires exactly one {key} value per rank, got {values}.") + if any(item <= 0 for item in values): + raise ValueError(f"Image DMD {key} must be positive, got {values}.") + return values[0] + + def encode_conditions( + self, + batch, + negative_prompt, + guidance_scale, + broadcast, + ): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + cached = _cached_condition_pair( + batch, + self.model, + require_negative=guidance_scale > 1, + ) + if cached is not None: + positive, negative = cached + return broadcast(positive), broadcast(negative) if negative is not None else None + with torch.no_grad(): + positive = self.model.encode_prompt_condition(prompt) + if guidance_scale > 1: + negative = self.model.encode_prompt_condition( + _negative_prompt( + conditioning, + negative_prompt, + scalar=scalar, + ) + ) + else: + negative = None + return ( + broadcast(positive), + broadcast(negative) if negative is not None else None, + ) + + def predict_velocity(self, latents, sigma, condition): + _require_singleton_tensor(latents, "DMD latent") + return self.model.predict_denoiser_output(latents, sigma, condition) + + def predict_guided_velocity( + self, + latents, + sigma, + condition, + negative_condition, + guidance_scale, + cfg_norm, + ): + if negative_condition is None: + return self.predict_velocity(latents, sigma, condition) + + if self._guidance_in_denoiser_space: + denoiser_input = self.model.prepare_denoiser_input(latents, condition=condition) + positive = self.model.denoise( + denoiser_input, + sigma, + condition, + ) + negative = self.model.denoise( + denoiser_input, + sigma, + negative_condition, + ) + prediction = self._cfg( + positive, + negative, + guidance_scale, + cfg_norm, + ) + return self.model.postprocess_denoiser_output( + prediction, + denoiser_input, + ) + return self._cfg( + self.predict_velocity(latents, sigma, condition), + self.predict_velocity(latents, sigma, negative_condition), + guidance_scale, + cfg_norm, + ) + + @staticmethod + def _cfg(positive, negative, scale, norm): + prediction = negative + scale * (positive - negative) + if norm in (None, "none"): + return prediction + if norm == "layer_norm": + positive_norm = torch.norm(positive, dim=-1, keepdim=True) + guided_norm = torch.norm(prediction, dim=-1, keepdim=True) + return prediction * (positive_norm / guided_norm.clamp_min(1e-12)) + if norm == "scalar": + ratio = torch.norm(positive) / torch.norm(prediction).clamp_min(1e-12) + return prediction * min(1.0, ratio.item()) + raise ValueError(f"Unsupported cfg_norm: {norm}") + + def initial_latents(self, latent_shape, dtype, broadcast): + if int(latent_shape[0]) != 1: + raise ValueError(f"DMD latent shape must start with 1, got {latent_shape}.") + return broadcast(torch.randn(latent_shape, device=self.device, dtype=dtype)) + + @staticmethod + def latent_hw(latent_shape): + return latent_shape[-2:] + + @staticmethod + def random_noise_like(latents, dtype, broadcast): + return broadcast(torch.randn_like(latents, dtype=dtype)) + + @staticmethod + def add_noise(scheduler, latents, noise, sigma): + return scheduler.add_noise(latents, noise, sigma) + + @staticmethod + def training_target(latents, noise): + return noise - latents.float() + + @staticmethod + def step(scheduler, velocity, step_index, sample): + return scheduler.step_by_index(velocity, step_index, sample) + + @staticmethod + def x0_from_velocity(sample, velocity, sigma): + if sigma.ndim == 0: + sigma = sigma.reshape(1) + expanded = sigma.reshape( + sigma.shape[0], + *([1] * (sample.ndim - 1)), + ) + return sample + (torch.zeros_like(expanded) - expanded) * velocity + + @staticmethod + def regression_loss(prediction, target): + return F.mse_loss( + prediction.float(), + target.float(), + reduction="mean", + ) + + @staticmethod + def dmd_loss(latents, fake_x0, teacher_x0): + with torch.no_grad(): + gradient = fake_x0 - teacher_x0 + dimensions = tuple(range(1, latents.ndim)) + normalizer = ( + (latents - teacher_x0) + .abs() + .mean( + dim=dimensions, + keepdim=True, + ) + ) + gradient = torch.nan_to_num(gradient / normalizer) + return 0.5 * F.mse_loss( + latents.float(), + (latents.float() - gradient.float()).detach(), + reduction="mean", + ) + + @staticmethod + def detach(value): + return value.detach() + + @staticmethod + def to_dtype(value, dtype): + return value.to(dtype=dtype) + + def extract_real_latents(self, batch, dtype, broadcast): + with torch.no_grad(): + latent = _cached_latent(batch, self.model) + if latent is None: + latent = self.model.encode_to_latent(batch) + latent = latent.to(device=self.device, dtype=dtype) + if latent.ndim == 4 and latent.shape[0] != 1: + latent = latent.unsqueeze(0) + _require_singleton_tensor(latent, "DMD real latent") + return broadcast(latent) + + def set_training(self, enabled: bool) -> None: + self.model.denoiser_module().train(enabled) + + def denoiser(self): + return self.model.denoiser_module() diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py new file mode 100644 index 000000000..382503435 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py @@ -0,0 +1,284 @@ +"""Composable consistency-model extensions for rectified-flow backbones.""" + +from __future__ import annotations + +import copy +from abc import ABC, abstractmethod +from collections.abc import Collection +from contextlib import contextmanager +from dataclasses import dataclass, replace + +import torch +from torch import Tensor, nn + +from .common import GenericConsistencyModelCapability + + +def _resolve_attribute(root, path: str): + value = root + for name in path.split("."): + try: + value = getattr(value, name) + except AttributeError as exc: + raise AttributeError(f"Cannot resolve {path!r}: {type(value).__name__} has no attribute {name!r}.") from exc + return value + + +def _module_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: + try: + parameter = next(module.parameters()) + except StopIteration as exc: + raise ValueError(f"{type(module).__name__} must contain parameters to define its device and dtype.") from exc + return parameter.device, parameter.dtype + + +class TimeEmbeddingAdapter(ABC): + """Describe how one denoiser creates its primary timestep embedding.""" + + @abstractmethod + def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: + """Return modules whose outputs receive endpoint conditioning.""" + + @abstractmethod + def base_embedder(self, denoiser: nn.Module) -> nn.Module: + """Return the train-time timestep MLP used by the pretrained model.""" + + def build_endpoint_embedder(self, denoiser: nn.Module) -> nn.Module: + return copy.deepcopy(self.base_embedder(denoiser)) + + @abstractmethod + def embedding_dimension(self, denoiser: nn.Module) -> int: + """Return the final timestep embedding width.""" + + @abstractmethod + def encode( + self, + denoiser: nn.Module, + embedder: nn.Module, + time: Tensor, + ) -> Tensor: + """Encode normalized flow time with ``embedder``.""" + + def add_endpoint_embedding( + self, + denoiser: nn.Module, + endpoint_embedder: nn.Module, + endpoint_time: Tensor, + hook_module: nn.Module, + output, + ): + del hook_module + if not torch.is_tensor(output): + raise TypeError(f"Endpoint conditioning requires a tensor timestep embedding, got {type(output).__name__}.") + embedding = self.encode(denoiser, endpoint_embedder, endpoint_time) + while embedding.ndim < output.ndim: + embedding = embedding.unsqueeze(-2) + return output + embedding.to(device=output.device, dtype=output.dtype) + + +@dataclass(frozen=True) +class ProjectedTimeEmbeddingAdapter(TimeEmbeddingAdapter): + """Adapter for backbones exposing separate time projection and MLP modules.""" + + hook_module_path: str + projection_module_path: str + embedding_module_path: str + embedding_dimension_path: str + time_scale: float = 1.0 + + def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: + return (_resolve_attribute(denoiser, self.hook_module_path),) + + def base_embedder(self, denoiser: nn.Module) -> nn.Module: + return _resolve_attribute(denoiser, self.embedding_module_path) + + def embedding_dimension(self, denoiser: nn.Module) -> int: + dimension = int(_resolve_attribute(denoiser, self.embedding_dimension_path)) + if dimension <= 0: + raise ValueError(f"Timestep embedding dimension must be positive, got {dimension}.") + return dimension + + def encode( + self, + denoiser: nn.Module, + embedder: nn.Module, + time: Tensor, + ) -> Tensor: + projection = _resolve_attribute(denoiser, self.projection_module_path) + device, dtype = _module_device_dtype(embedder) + scaled_time = time.to(device=device) + if self.time_scale != 1.0: + scaled_time = scaled_time * self.time_scale + projected = projection(scaled_time) + return embedder(projected.to(device=device, dtype=dtype)) + + +@dataclass(frozen=True) +class SinusoidalTimeEmbeddingAdapter(TimeEmbeddingAdapter): + """Adapter for Wan-style MLPs whose sinusoidal projection lives in forward.""" + + embedding_module_path: str + embedding_dimension_path: str + frequency_dimension_path: str + time_scale: float = 1.0 + + def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: + return (self.base_embedder(denoiser),) + + def base_embedder(self, denoiser: nn.Module) -> nn.Module: + return _resolve_attribute(denoiser, self.embedding_module_path) + + def embedding_dimension(self, denoiser: nn.Module) -> int: + dimension = int(_resolve_attribute(denoiser, self.embedding_dimension_path)) + if dimension <= 0: + raise ValueError(f"Timestep embedding dimension must be positive, got {dimension}.") + return dimension + + def encode( + self, + denoiser: nn.Module, + embedder: nn.Module, + time: Tensor, + ) -> Tensor: + frequency_dimension = int(_resolve_attribute(denoiser, self.frequency_dimension_path)) + if frequency_dimension <= 0 or frequency_dimension % 2: + raise ValueError(f"Sinusoidal frequency dimension must be a positive even integer, got {frequency_dimension}.") + + device, dtype = _module_device_dtype(embedder) + position = time.to(device=device).reshape(-1) + if self.time_scale != 1.0: + position = position * self.time_scale + projection = self._sinusoidal_embedding(frequency_dimension, position) + return embedder(projection.to(device=device, dtype=dtype)) + + @staticmethod + def _sinusoidal_embedding(dimension: int, position: Tensor) -> Tensor: + half = dimension // 2 + position = position.to(torch.float64) + frequencies = torch.pow( + 10000, + -torch.arange(half, device=position.device, dtype=torch.float64) / half, + ) + phases = torch.outer(position, frequencies) + return torch.cat((torch.cos(phases), torch.sin(phases)), dim=1) + + +class TimeConditionedConsistencyModelCapability(GenericConsistencyModelCapability): + """Generic CM adapter with optional endpoint conditioning and loss variance. + + The capability owns objective-specific state and modules. The model wrapper + only supplies a small, immutable description of its timestep embedding. + """ + + _SUPPORTED_FEATURES = frozenset({"endpoint_time", "log_variance"}) + + def __init__( + self, + model, + time_adapter: TimeEmbeddingAdapter, + *, + endpoint_module_name: str = "consistency_endpoint_embedder", + log_variance_module_name: str = "consistency_logvar_head", + ) -> None: + super().__init__(model) + self._time_adapter = time_adapter + self._endpoint_module_name = endpoint_module_name + self._log_variance_module_name = log_variance_module_name + self._configured_features: set[str] = set() + self._endpoint_time: Tensor | None = None + self._endpoint_hooks = [] + + def configure(self, features: Collection[str]) -> None: + features = frozenset(features) + unsupported = features - self._SUPPORTED_FEATURES + if unsupported: + names = ", ".join(sorted(unsupported)) + raise NotImplementedError(f"{type(self).__name__} does not support consistency features: {names}.") + + denoiser = self.denoiser() + if "log_variance" in features and not hasattr(denoiser, self._log_variance_module_name): + source_embedder = self._time_adapter.base_embedder(denoiser) + device, dtype = _module_device_dtype(source_embedder) + head = nn.Linear( + self._time_adapter.embedding_dimension(denoiser), + 1, + device=device, + dtype=dtype, + ) + setattr(denoiser, self._log_variance_module_name, head) + + if "endpoint_time" in features: + if not hasattr(denoiser, self._endpoint_module_name): + endpoint_embedder = self._time_adapter.build_endpoint_embedder(denoiser) + setattr(denoiser, self._endpoint_module_name, endpoint_embedder) + if not self._endpoint_hooks: + self._endpoint_hooks = [module.register_forward_hook(self._add_endpoint_time_embedding) for module in self._time_adapter.hook_modules(denoiser)] + + self._configured_features.update(features) + + def restore_trainable_auxiliary(self) -> None: + denoiser = self.denoiser() + for module_name in self._configured_module_names(): + getattr(denoiser, module_name).requires_grad_(True) + + def auxiliary_parameter_names(self) -> tuple[str, ...]: + module_names = self._configured_module_names() + return tuple(name for name, _ in self.denoiser().named_parameters() if name.partition(".")[0] in module_names) + + def predict(self, request, path): + with self._configured_prediction(request) as prepared_request: + return super().predict(prepared_request, path) + + @contextmanager + def _configured_prediction(self, request): + model_kwargs = dict(request.model_kwargs) + endpoint_time = model_kwargs.pop("endpoint_time", None) + if endpoint_time is not None and "endpoint_time" not in self._configured_features: + raise RuntimeError("endpoint_time was provided before endpoint conditioning was configured.") + + request = replace(request, model_kwargs=model_kwargs) + with self._use_endpoint_time(endpoint_time): + yield request + + def predict_log_variance(self, time: Tensor) -> Tensor: + if "log_variance" not in self._configured_features: + raise RuntimeError("The consistency log-variance head has not been configured.") + denoiser = self.denoiser() + head = getattr(denoiser, self._log_variance_module_name) + embedding = self._time_adapter.encode( + denoiser, + self._time_adapter.base_embedder(denoiser), + time, + ) + return head(embedding.to(device=head.weight.device, dtype=head.weight.dtype)) + + def _configured_module_names(self) -> set[str]: + names = set() + if "endpoint_time" in self._configured_features: + names.add(self._endpoint_module_name) + if "log_variance" in self._configured_features: + names.add(self._log_variance_module_name) + return names + + @contextmanager + def _use_endpoint_time(self, endpoint_time): + previous = self._endpoint_time + self._endpoint_time = endpoint_time + try: + yield + finally: + self._endpoint_time = previous + + def _add_endpoint_time_embedding(self, module, inputs, output): + del inputs + if self._endpoint_time is None: + return output + denoiser = self.denoiser() + endpoint_embedder = getattr(denoiser, self._endpoint_module_name) + return self._time_adapter.add_endpoint_embedding( + denoiser, + endpoint_embedder, + self._endpoint_time, + module, + output, + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/latent_geometry.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/latent_geometry.py new file mode 100644 index 000000000..65c64caee --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/latent_geometry.py @@ -0,0 +1,71 @@ +"""Model-agnostic descriptions of latent tensor geometry.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Protocol + + +def _resolve_attribute(root, path: str): + value = root + for name in path.split("."): + try: + value = getattr(value, name) + except AttributeError as exc: + raise AttributeError(f"Cannot resolve {path!r}: {type(value).__name__} has no attribute {name!r}.") from exc + return value + + +class LatentGeometry(Protocol): + """Build a model's latent tensor shape from an output spatial size.""" + + def shape(self, model, height: int, width: int) -> tuple[int, ...]: + """Return a singleton latent shape for ``height`` x ``width`` output.""" + + +@dataclass(frozen=True) +class SpatialLatentGeometry: + """Declarative geometry for image-like 4D or single-frame 5D latents. + + ``spatial_downsample_multiplier`` describes additional downsampling on top + of the model's VAE scale. ``input_size_multiplier`` captures packing + requirements such as an even VAE-latent grid. + """ + + channels_path: str + spatial_scale_path: str = "vae_scale_factor" + spatial_downsample_multiplier: int = 1 + input_size_multiplier: int = 2 + temporal_size: int | None = None + + def __post_init__(self) -> None: + if self.spatial_downsample_multiplier <= 0: + raise ValueError("spatial_downsample_multiplier must be positive.") + if self.input_size_multiplier <= 0: + raise ValueError("input_size_multiplier must be positive.") + if self.temporal_size is not None and self.temporal_size <= 0: + raise ValueError("temporal_size must be positive when provided.") + + def shape(self, model, height: int, width: int) -> tuple[int, ...]: + height, width = int(height), int(width) + channels = int(_resolve_attribute(model, self.channels_path)) + spatial_scale = int(_resolve_attribute(model, self.spatial_scale_path)) + if channels <= 0 or spatial_scale <= 0: + raise ValueError(f"Latent channels and spatial scale must be positive, got channels={channels}, scale={spatial_scale}.") + + spatial_downsample = spatial_scale * self.spatial_downsample_multiplier + required_multiple = math.lcm( + spatial_downsample, + spatial_scale * self.input_size_multiplier, + ) + if height <= 0 or width <= 0 or height % required_multiple or width % required_multiple: + raise ValueError(f"Output height and width must be positive multiples of {required_multiple}, got {height}x{width}.") + + spatial_shape = ( + height // spatial_downsample, + width // spatial_downsample, + ) + if self.temporal_size is None: + return 1, channels, *spatial_shape + return 1, channels, self.temporal_size, *spatial_shape diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/__init__.py new file mode 100644 index 000000000..27b4d5f6d --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/__init__.py @@ -0,0 +1 @@ +"""Flux2 model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py new file mode 100644 index 000000000..010fd2d89 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -0,0 +1,9 @@ +"""Capabilities implemented specifically by Flux2 models.""" + +from .flux2_consistency_model_capability import Flux2ConsistencyModelCapability +from .flux2_edit_distribution_matching_capability import Flux2EditDistributionMatchingCapability + +__all__ = [ + "Flux2ConsistencyModelCapability", + "Flux2EditDistributionMatchingCapability", +] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py new file mode 100644 index 000000000..faf3281db --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py @@ -0,0 +1,22 @@ +"""Consistency-model capability for Flux2 backbones.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( + ProjectedTimeEmbeddingAdapter, + TimeConditionedConsistencyModelCapability, +) + + +class Flux2ConsistencyModelCapability(TimeConditionedConsistencyModelCapability): + """Bind generic consistency extensions to Flux2 timestep conditioning.""" + + def __init__(self, model) -> None: + super().__init__( + model, + ProjectedTimeEmbeddingAdapter( + hook_module_path="time_guidance_embed", + projection_module_path="time_guidance_embed.time_proj", + embedding_module_path="time_guidance_embed.timestep_embedder", + embedding_dimension_path="inner_dim", + time_scale=1000.0, + ), + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_edit_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_edit_distribution_matching_capability.py new file mode 100644 index 000000000..31a357604 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_edit_distribution_matching_capability.py @@ -0,0 +1,32 @@ +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _cached_condition_pair, + _negative_prompt, + _require_single_prompt, +) + + +class Flux2EditDistributionMatchingCapability(GenericDistributionMatchingCapability): + """Keep reference-image conditioning in Flux2 edit DMD.""" + + cache_uses_sample_context = True + + def encode_conditions(self, batch, negative_prompt, guidance_scale, broadcast): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + cached = _cached_condition_pair(batch, self.model, require_negative=guidance_scale > 1) + if cached is not None: + positive, negative = cached + return broadcast(positive), broadcast(negative) if negative is not None else None + prompts = [prompt] + if guidance_scale > 1: + prompts.append(_negative_prompt(conditioning, negative_prompt, scalar=scalar)) + + with torch.no_grad(): + conditions = self.model.encode_conditions_with_source(batch, prompts) + positive = conditions[0] + negative = conditions[1] if len(conditions) > 1 else None + return broadcast(positive), broadcast(negative) if negative is not None else None diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py new file mode 100644 index 000000000..323fb124e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py @@ -0,0 +1,257 @@ +from dataclasses import dataclass + +import torch +from diffusers import AutoencoderKLFlux2, Flux2Transformer2DModel +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor + +from ..base import BaseModel + + +@dataclass +class Flux2DenoiserInput: + hidden_states: torch.Tensor + image_ids: torch.Tensor + target_ids: torch.Tensor + target_token_length: int + height: int + width: int + + +class Flux2ModelBase(BaseModel): + """Shared tensor mechanics for the Flux2 model family.""" + + requires_source_images = False + target_latent_mode = "sample" + default_text_encoder_out_layers = () + default_unconditional_prompt = "" + + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + model_path = self.config["model"]["pretrained_model_name_or_path"] + self._validate_model_path(model_path) + self._configure_model() + + if load_condition_encoder: + self._load_condition_encoder(model_path) + if load_vae: + self._load_vae(model_path) + else: + self._load_vae_config(model_path) + if load_transformer: + self.transformer = self.load_transformer(model_path) + + def _load_condition_encoder(self, model_path): + self.text_pipeline = self.pipeline_cls.from_pretrained( + model_path, + transformer=None, + vae=None, + torch_dtype=self.running_dtype, + ).to(self.device) + self.text_pipeline.text_encoder.requires_grad_(False) + self.text_pipeline.text_encoder.eval() + + def _load_vae(self, model_path): + self.vae = AutoencoderKLFlux2.from_pretrained( + model_path, + subfolder="vae", + torch_dtype=self.running_dtype, + ).to(self.device) + self.vae_config = self.vae.config + self.vae.requires_grad_(False) + self.vae.eval() + self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + + def _load_vae_config(self, model_path): + self.vae_config = AutoencoderKLFlux2.load_config(model_path, subfolder="vae") + + def _validate_model_path(self, model_path): + del model_path + + def _configure_model(self): + pass + + def reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self._reuse_model_state_from(source) + + def _reuse_model_state_from(self, source): + del source + + def load_transformer(self, model_path=None): + model_path = model_path or self.config["model"]["pretrained_model_name_or_path"] + return Flux2Transformer2DModel.from_pretrained( + model_path, + subfolder="transformer", + torch_dtype=self.running_dtype, + ).to(self.device) + + def denoiser_module(self): + return self.transformer + + def fsdp2_shard_plan(self, fsdp_config): + reshard_config = fsdp_config["reshard_after_forward"] + return [ + { + "modules": self.transformer.transformer_blocks, + "reshard_after_forward": reshard_config["block_reshard"], + }, + { + "modules": self.transformer.single_transformer_blocks, + "reshard_after_forward": reshard_config["block_reshard"], + }, + { + "module": self.transformer, + "reshard_after_forward": reshard_config["root_reshard"], + }, + ] + + @property + def vae_scale_factor(self): + config = self.vae_config + block_out_channels = config["block_out_channels"] if isinstance(config, dict) else config.block_out_channels + return 2 ** (len(block_out_channels) - 1) + + def _normalize_patch_latents(self, latents): + latents = self.pipeline_cls._patchify_latents(latents) + mean, std = self._latent_statistics(latents) + return (latents - mean) / std + + def _denormalize_patch_latents(self, latents): + mean, std = self._latent_statistics(latents) + return self.pipeline_cls._unpatchify_latents(latents * std + mean) + + def _latent_statistics(self, latents): + mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + variance = self.vae.bn.running_var.view(1, -1, 1, 1).to(latents.device, latents.dtype) + return mean, torch.sqrt(variance + self.vae.config.batch_norm_eps) + + def encode_to_latent(self, sample): + return self._encode_target_latent(sample, mode=self.target_latent_mode) + + def encode_to_cache_latent(self, sample): + return self._encode_target_latent(sample, mode="mode") + + def _encode_target_latent(self, sample, *, mode): + image = sample["inputs"]["target_pixel_values"] + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4: + raise ValueError(f"Expected target_pixel_values with shape [B, C, H, W], got {tuple(image.shape)}") + image = image.to(device=self.device, dtype=self.running_dtype) + distribution = self.vae.encode(image).latent_dist + latent = getattr(distribution, mode)() + return self._normalize_patch_latents(latent) + + def encode_condition(self, sample): + prompt = sample["conditioning"]["prompt"] + if self.requires_source_images: + return self.encode_conditions_with_source(sample, [prompt])[0] + return self.encode_prompt_condition(prompt) + + def encode_conditions_with_context(self, sample, prompts): + if self.requires_source_images: + return self.encode_conditions_with_source(sample, prompts) + return super().encode_conditions_with_context(sample, prompts) + + def encode_prompt_condition(self, prompt): + model_config = self.config["model"] + prompt_embed, text_ids = self.text_pipeline.encode_prompt( + prompt=prompt, + device=self.device, + num_images_per_prompt=1, + max_sequence_length=model_config.get("max_sequence_length", 512), + text_encoder_out_layers=tuple(model_config.get("text_encoder_out_layers", self.default_text_encoder_out_layers)), + ) + return {"prompt_embed": prompt_embed, "text_ids": text_ids} + + def encode_conditions_with_source(self, sample, prompts): + images = sample["inputs"].get("source_vae_pixel_values", []) + if not images: + raise ValueError(f"{type(self).__name__} requires at least one source image") + + reference_tokens, reference_ids = self._encode_reference_images(images) + conditions = [] + for prompt in prompts: + condition = self.encode_prompt_condition(prompt) + condition.update(reference_tokens=reference_tokens, reference_ids=reference_ids) + conditions.append(condition) + return conditions + + def _encode_reference_images(self, images): + latents = [] + for image in images: + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4: + raise ValueError(f"Expected source image with shape [B, C, H, W], got {tuple(image.shape)}") + image = image.to(device=self.device, dtype=self.running_dtype) + latent = self.vae.encode(image).latent_dist.mode() + latents.append(self._normalize_patch_latents(latent)) + + batch_size = latents[0].shape[0] + if any(latent.shape[0] != batch_size for latent in latents): + raise ValueError("All source images must have the same batch size") + + reference_ids = self.pipeline_cls._prepare_image_ids([latent[:1] for latent in latents]) + reference_ids = reference_ids.expand(batch_size, -1, -1).to(self.device) + reference_tokens = torch.cat([self.pipeline_cls._pack_latents(latent) for latent in latents], dim=1) + return reference_tokens, reference_ids + + def prepare_denoiser_input(self, noisy_latent, condition=None): + height, width = noisy_latent.shape[-2:] + target_tokens = self.pipeline_cls._pack_latents(noisy_latent) + target_ids = self.pipeline_cls._prepare_latent_ids(noisy_latent).to(noisy_latent.device) + + hidden_states = target_tokens + image_ids = target_ids + if condition is not None and condition.get("reference_tokens") is not None: + reference_tokens = condition["reference_tokens"] + reference_ids = condition.get("reference_ids") + if reference_ids is None: + raise ValueError("Flux2 reference tokens require matching reference IDs") + if reference_tokens.shape[0] != target_tokens.shape[0]: + raise ValueError("Flux2 target and reference tokens must have the same batch size") + reference_tokens = reference_tokens.to(target_tokens.device, target_tokens.dtype) + reference_ids = reference_ids.to(target_ids.device) + hidden_states = torch.cat([target_tokens, reference_tokens], dim=1) + image_ids = torch.cat([target_ids, reference_ids], dim=1) + + return Flux2DenoiserInput( + hidden_states=hidden_states, + image_ids=image_ids, + target_ids=target_ids, + target_token_length=target_tokens.shape[1], + height=height, + width=width, + ) + + def _denoise(self, denoiser_input, timestep_or_sigma, condition, guidance): + prediction = self.transformer( + hidden_states=denoiser_input.hidden_states, + timestep=timestep_or_sigma, + guidance=guidance, + encoder_hidden_states=condition["prompt_embed"], + txt_ids=condition["text_ids"], + img_ids=denoiser_input.image_ids, + joint_attention_kwargs={}, + return_dict=False, + )[0] + return prediction[:, : denoiser_input.target_token_length] + + def postprocess_denoiser_output(self, prediction, denoiser_input): + return self.pipeline_cls._unpack_latents_with_ids(prediction, denoiser_input.target_ids) + + def prepare_infer_latents(self, height, width, generator=None): + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + shape = (1, self.transformer.config.in_channels, latent_height // 2, latent_width // 2) + return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) + + def decode_latent(self, latent): + image = self.vae.decode(self._denormalize_patch_latents(latent)).sample + return self.image_processor.postprocess(image, output_type="pil") diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py new file mode 100644 index 000000000..ba3b825eb --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -0,0 +1,164 @@ +import torch +from diffusers import AutoencoderKLFlux2 +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor + +from lightx2v_train.utils.image_ops import align_dimension, calculate_area_dimensions, resize_and_center_crop +from lightx2v_train.utils.registry import SAMPLE_PROCESSOR_REGISTER + + +def _size_multiple_from_config(config): + processor_config = config.get("data", {}).get("processor", {}) + value = processor_config.get("size_multiple") + if value is None: + model_path = config["model"]["pretrained_model_name_or_path"] + vae_config = AutoencoderKLFlux2.load_config(model_path, subfolder="vae") + value = 2 ** (len(vae_config["block_out_channels"]) - 1) * 2 + value = int(value) + if value <= 0: + raise ValueError(f"size_multiple must be positive, got {value}") + return value + + +def _target_area_from_config(config): + processor_config = config.get("data", {}).get("processor", {}) + return int(processor_config.get("target_area", 1024 * 1024)) + + +def _optional_scalar(mapping, key): + value = mapping.get(key) + if value is None: + return None + if torch.is_tensor(value): + if value.numel() != 1: + raise ValueError(f"{key} must contain one value, got {value.numel()}") + value = value.item() + elif isinstance(value, (list, tuple)): + if len(value) != 1: + raise ValueError(f"{key} must contain one value, got {len(value)}") + value = value[0] + if torch.is_tensor(value): + value = value.item() + return int(value) + + +def _explicit_target_size(sample): + meta = sample.get("meta", {}) + height = _optional_scalar(meta, "target_height") + width = _optional_scalar(meta, "target_width") + if (height is None) != (width is None): + raise ValueError("meta.target_height and meta.target_width must be provided together") + return height, width + + +@SAMPLE_PROCESSOR_REGISTER("flux2_dev") +@SAMPLE_PROCESSOR_REGISTER("flux2_klein") +class Flux2DataProcessor: + def __init__(self, config): + self.image_processor = Flux2ImageProcessor(vae_scale_factor=_size_multiple_from_config(config)) + self.target_area = _target_area_from_config(config) + self.unconditional_prompt = config["model"].get("unconditional_prompt", "") + + def __call__(self, sample): + inputs = sample["inputs"] + image = inputs.pop("target_image", None) + if image is not None: + inputs["target_pixel_values"] = self._process_target(image, sample) + return sample + + def _process_target(self, image, sample, fallback_size=None): + height, width = _explicit_target_size(sample) + multiple = int(self.image_processor.config.vae_scale_factor) + if height is not None: + height = align_dimension(height, multiple) + width = align_dimension(width, multiple) + elif fallback_size is not None: + height, width = fallback_size + else: + width, height = calculate_area_dimensions( + self.target_area, + image.width / image.height, + multiple, + ) + image = resize_and_center_crop(image, width, height) + return self.image_processor.preprocess(image)[0] + + def infer_target_size(self, sample, default_height, default_width): + height, width = _explicit_target_size(sample) + height = default_height if height is None else height + width = default_width if width is None else width + multiple = int(self.image_processor.config.vae_scale_factor) + return align_dimension(int(height), multiple), align_dimension(int(width), multiple) + + +@SAMPLE_PROCESSOR_REGISTER("flux2_dev_edit") +@SAMPLE_PROCESSOR_REGISTER("flux2_klein_edit") +class Flux2EditDataProcessor(Flux2DataProcessor): + def __call__(self, sample): + inputs = sample["inputs"] + meta = sample["meta"] + target_image = inputs.pop("target_image", None) + source_images = inputs.pop("source_images", []) + if not source_images: + raise ValueError("Flux2 image editing requires at least one source image") + + source_pixels = [] + source_sizes = [] + for image in source_images: + pixels, size = self._process_source(image) + source_pixels.append(pixels) + source_sizes.append(size) + + reference_height, reference_width = source_sizes[0] + meta["reference_image_height"] = reference_height + meta["reference_image_width"] = reference_width + inputs["source_vae_pixel_values"] = source_pixels + if target_image is not None: + target_pixels = self._process_target( + target_image, + sample, + fallback_size=(reference_height, reference_width), + ) + inputs["target_pixel_values"] = target_pixels + meta["target_height"] = int(target_pixels.shape[-2]) + meta["target_width"] = int(target_pixels.shape[-1]) + else: + target_height, target_width = _explicit_target_size(sample) + if target_height is None: + target_height, target_width = reference_height, reference_width + else: + multiple = int(self.image_processor.config.vae_scale_factor) + target_height = align_dimension(target_height, multiple) + target_width = align_dimension(target_width, multiple) + meta["target_height"] = target_height + meta["target_width"] = target_width + return sample + + def _process_source(self, image): + self.image_processor.check_image_input(image) + if image.width * image.height > self.target_area: + image = self.image_processor._resize_to_target_area(image, self.target_area) + + multiple = int(self.image_processor.config.vae_scale_factor) + width = align_dimension(image.width, multiple) + height = align_dimension(image.height, multiple) + pixels = self.image_processor.preprocess( + image, + height=height, + width=width, + resize_mode="crop", + )[0] + return pixels, (height, width) + + def infer_target_size(self, sample, default_height, default_width): + del default_height, default_width + height, width = _explicit_target_size(sample) + if height is not None: + multiple = int(self.image_processor.config.vae_scale_factor) + return align_dimension(height, multiple), align_dimension(width, multiple) + + meta = sample.get("meta", {}) + height = _optional_scalar(meta, "reference_image_height") + width = _optional_scalar(meta, "reference_image_width") + if height is None or width is None: + raise ValueError("Flux2 image editing requires source-derived target dimensions") + return height, width diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py new file mode 100644 index 000000000..81414648b --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -0,0 +1,56 @@ +import torch +from diffusers import Flux2Pipeline + +from lightx2v_train.model_capabilities import ConsistencyModelCapability, DistributionMatchingCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyModelCapability +from lightx2v_train.utils.registry import MODEL_REGISTER + +from .common import Flux2ModelBase + + +@MODEL_REGISTER("flux2_dev") +class Flux2DevModel(Flux2ModelBase): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/black-forest-labs/FLUX.2-dev + """ + + pipeline_cls = Flux2Pipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability + default_text_encoder_out_layers = (10, 20, 30) + + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register(FlowMatchingSFTCapability, GenericFlowMatchingCapability(self)) + self.capabilities.register( + DistributionMatchingCapability, + self.distribution_matching_capability_cls( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="transformer.config.in_channels", + spatial_downsample_multiplier=2, + ), + guidance_in_denoiser_space=True, + ), + ) + self.capabilities.register(ConsistencyModelCapability, Flux2ConsistencyModelCapability(self)) + + def _configure_model(self): + self.guidance_scale = float(self.config["model"].get("guidance_scale", 4.0)) + + def _reuse_model_state_from(self, source): + self.guidance_scale = source.guidance_scale + + def denoise(self, denoiser_input, timestep_or_sigma, condition): + guidance = torch.full( + (denoiser_input.hidden_states.shape[0],), + self.guidance_scale, + device=self.device, + dtype=torch.float32, + ) + return self._denoise(denoiser_input, timestep_or_sigma, condition, guidance) + + def apply_cfg(self, positive, negative, guidance_scale): + del positive, negative, guidance_scale + raise ValueError("Flux2 Dev uses embedded guidance and does not support external classifier-free guidance") diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py new file mode 100644 index 000000000..04598a17b --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py @@ -0,0 +1,16 @@ +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2EditDistributionMatchingCapability +from lightx2v_train.utils.registry import MODEL_REGISTER + +from .flux2_dev import Flux2DevModel + + +@MODEL_REGISTER("flux2_dev_edit") +class Flux2DevEditModel(Flux2DevModel): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/black-forest-labs/FLUX.2-dev + """ + + requires_source_images = True + target_latent_mode = "mode" + distribution_matching_capability_cls = Flux2EditDistributionMatchingCapability + shared_condition_keys = ("reference_tokens", "reference_ids") diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py new file mode 100644 index 000000000..f52ff1fa2 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -0,0 +1,59 @@ +from diffusers import Flux2KleinPipeline + +from lightx2v_train.model_capabilities import ( + ConsistencyModelCapability, + DistributionMatchingCapability, + FlowMatchingSFTCapability, +) +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.flux2.capability_adapters import ( + Flux2ConsistencyModelCapability, +) +from lightx2v_train.utils.registry import MODEL_REGISTER + +from .common import Flux2ModelBase + + +@MODEL_REGISTER("flux2_klein") +class Flux2KleinModel(Flux2ModelBase): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B + - https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B + """ + + pipeline_cls = Flux2KleinPipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability + default_text_encoder_out_layers = (9, 18, 27) + + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register(FlowMatchingSFTCapability, GenericFlowMatchingCapability(self)) + self.capabilities.register( + DistributionMatchingCapability, + self.distribution_matching_capability_cls( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="transformer.config.in_channels", + spatial_downsample_multiplier=2, + ), + guidance_in_denoiser_space=True, + ), + ) + self.capabilities.register(ConsistencyModelCapability, Flux2ConsistencyModelCapability(self)) + + def _validate_model_path(self, model_path): + pipeline_config = self.pipeline_cls.load_config(model_path) + if pipeline_config.get("is_distilled", False): + raise ValueError("Flux2KleinModel supports only FLUX.2-klein-base-4B/9B checkpoints; step-distilled FLUX.2-klein-4B/9B checkpoints are unsupported") + + def denoise(self, denoiser_input, timestep_or_sigma, condition): + return self._denoise(denoiser_input, timestep_or_sigma, condition, guidance=None) + + def postprocess_denoiser_output(self, prediction, denoiser_input): + return self.pipeline_cls._unpack_latents_with_ids( + prediction, + denoiser_input.target_ids, + height=denoiser_input.height, + width=denoiser_input.width, + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py new file mode 100644 index 000000000..af84df2ef --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py @@ -0,0 +1,17 @@ +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2EditDistributionMatchingCapability +from lightx2v_train.utils.registry import MODEL_REGISTER + +from .flux2_klein import Flux2KleinModel + + +@MODEL_REGISTER("flux2_klein_edit") +class Flux2KleinEditModel(Flux2KleinModel): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B + - https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B + """ + + requires_source_images = True + target_latent_mode = "mode" + distribution_matching_capability_cls = Flux2EditDistributionMatchingCapability + shared_condition_keys = ("reference_tokens", "reference_ids") diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py deleted file mode 100644 index c140f1ad4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py +++ /dev/null @@ -1,179 +0,0 @@ -from dataclasses import dataclass - -import torch -from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel -from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor - -from lightx2v_train.utils.registry import MODEL_REGISTER - -from .base import BaseModel - - -@dataclass -class Flux2DevDenoiserInput: - hidden_states: torch.Tensor - img_ids: torch.Tensor - height: int - width: int - - -@MODEL_REGISTER("flux2_dev") -class Flux2DevModel(BaseModel): - pipeline_cls = Flux2Pipeline - - def load_components(self, transformer_only=False, reference_model=None): - model_config = self.config["model"] - model_path = model_config["pretrained_model_name_or_path"] - self.guidance_scale = float(model_config.get("guidance_scale", 4.0)) - - if transformer_only: - if reference_model is not None: - self.text_pipeline = reference_model.text_pipeline - self.vae = reference_model.vae - self.image_processor = reference_model.image_processor - self.guidance_scale = reference_model.guidance_scale - self.transformer = self.load_transformer() - return - - self.text_pipeline = Flux2Pipeline.from_pretrained( - model_path, - transformer=None, - vae=None, - torch_dtype=self.running_dtype, - ).to(self.device) - self.vae = AutoencoderKLFlux2.from_pretrained(model_path, subfolder="vae", torch_dtype=self.running_dtype).to(self.device) - self.transformer = self.load_transformer() - - self.text_pipeline.text_encoder.requires_grad_(False) - self.vae.requires_grad_(False) - self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - - def load_transformer(self): - model_path = self.config["model"]["pretrained_model_name_or_path"] - return Flux2Transformer2DModel.from_pretrained(model_path, subfolder="transformer", torch_dtype=self.running_dtype).to(self.device) - - def denoiser_module(self): - return self.transformer - - def fsdp2_shard_plan(self, fsdp_config): - reshard_config = fsdp_config["reshard_after_forward"] - return [ - { - "modules": self.transformer.transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "modules": self.transformer.single_transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "module": self.transformer, - "reshard_after_forward": reshard_config["root_reshard"], - }, - ] - - @property - def vae_scale_factor(self): - return 2 ** (len(self.vae.config.block_out_channels) - 1) - - def _normalize_patch_latents(self, latents): - latents = Flux2Pipeline._patchify_latents(latents) - latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) - latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(latents.device, latents.dtype) - return (latents - latents_bn_mean) / latents_bn_std - - def _denormalize_patch_latents(self, latents): - latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) - latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(latents.device, latents.dtype) - latents = latents * latents_bn_std + latents_bn_mean - return Flux2Pipeline._unpatchify_latents(latents) - - def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) - latent = self.vae.encode(image).latent_dist.sample() - return self._normalize_patch_latents(latent) - - def encode_condition(self, sample): - prompt = sample["conditioning"]["prompt"] - return self.encode_prompt_condition(prompt) - - def encode_prompt_condition(self, prompt): - model_config = self.config["model"] - prompt_embed, text_ids = self.text_pipeline.encode_prompt( - prompt=prompt, - device=self.device, - num_images_per_prompt=1, - max_sequence_length=model_config.get("max_sequence_length", 512), - text_encoder_out_layers=tuple(model_config.get("text_encoder_out_layers", (10, 20, 30))), - ) - return {"prompt_embed": prompt_embed, "text_ids": text_ids} - - def prepare_denoiser_input(self, noisy_latent, condition=None): - h, w = noisy_latent.shape[2], noisy_latent.shape[3] - packed = Flux2Pipeline._pack_latents(noisy_latent) - img_ids = Flux2Pipeline._prepare_latent_ids(noisy_latent).to(self.device) - return Flux2DevDenoiserInput( - hidden_states=packed, - img_ids=img_ids, - height=h, - width=w, - ) - - def denoise(self, denoiser_input, timestep_or_sigma, condition): - guidance = torch.full( - (denoiser_input.hidden_states.shape[0],), - self.guidance_scale, - device=self.device, - dtype=torch.float32, - ) - return self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, - guidance=guidance, - encoder_hidden_states=condition["prompt_embed"], - txt_ids=condition["text_ids"], - img_ids=denoiser_input.img_ids, - joint_attention_kwargs={}, - return_dict=False, - )[0] - - def postprocess_denoiser_output(self, prediction, denoiser_input): - return Flux2Pipeline._unpack_latents_with_ids(prediction, denoiser_input.img_ids) - - def prepare_infer_latents(self, height, width, generator=None): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - shape = (1, self.transformer.config.in_channels, latent_h // 2, latent_w // 2) - return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) - - def dmd_latent_shape(self, batch_size, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return (int(batch_size), int(self.transformer.config.in_channels), latent_h // 2, latent_w // 2) - - def cfg_on_denoiser_output(self): - return True - - def decode_latent(self, latent): - latent = self._denormalize_patch_latents(latent) - image = self.vae.decode(latent).sample - return self.image_processor.postprocess(image, output_type="pil") - - def assemble_pipeline(self, scheduler=None): - return Flux2Pipeline( - tokenizer=self.text_pipeline.tokenizer, - text_encoder=self.text_pipeline.text_encoder, - vae=self.vae, - transformer=self.transformer, - scheduler=scheduler or self.text_pipeline.scheduler, - ).to(self.device) - - def get_pipeline_infer_kwargs(self, infer_config): - return { - "height": infer_config.get("height", infer_config.get("default_height", 1024)), - "width": infer_config.get("width", infer_config.get("default_width", 1024)), - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", self.guidance_scale), - "max_sequence_length": self.config["model"].get("max_sequence_length", 512), - "text_encoder_out_layers": tuple(self.config["model"].get("text_encoder_out_layers", (10, 20, 30))), - } diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py deleted file mode 100644 index c638c476b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py +++ /dev/null @@ -1,231 +0,0 @@ -from dataclasses import dataclass - -import torch -from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel -from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor - -from lightx2v_train.utils.registry import MODEL_REGISTER - -from .base import BaseModel - - -@dataclass -class Flux2KleinDenoiserInput: - hidden_states: torch.Tensor - img_ids: torch.Tensor - height: int - width: int - - -@MODEL_REGISTER("flux2_klein") -class Flux2KleinModel(BaseModel): - pipeline_cls = Flux2KleinPipeline - - def load_components(self, transformer_only=False, reference_model=None): - if transformer_only: - if reference_model is not None: - self.text_pipeline = reference_model.text_pipeline - self.vae = reference_model.vae - self.image_processor = reference_model.image_processor - self.transformer = self.load_transformer() - return - model_path = self.config["model"]["pretrained_model_name_or_path"] - self.text_pipeline = Flux2KleinPipeline.from_pretrained( - model_path, - transformer=None, - vae=None, - torch_dtype=self.running_dtype, - ).to(self.device) - self.vae = AutoencoderKLFlux2.from_pretrained(model_path, subfolder="vae").to(self.device, dtype=self.running_dtype) - self.transformer = self.load_transformer() - - self.text_pipeline.text_encoder.requires_grad_(False) - self.vae.requires_grad_(False) - self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - - def load_transformer(self, model_path=None): - model_path = model_path or self.config["model"]["pretrained_model_name_or_path"] - return Flux2Transformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) - - def denoiser_module(self): - return self.transformer - - def fsdp2_shard_plan(self, fsdp_config): - reshard_config = fsdp_config["reshard_after_forward"] - return [ - { - "modules": self.transformer.transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "modules": self.transformer.single_transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "module": self.transformer, - "reshard_after_forward": reshard_config["root_reshard"], - }, - ] - - @property - def vae_scale_factor(self): - return 2 ** (len(self.vae.config.block_out_channels) - 1) - - def _normalize_patch_latents(self, latents): - latents = Flux2KleinPipeline._patchify_latents(latents) - latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) - latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(latents.device, latents.dtype) - return (latents - latents_bn_mean) / latents_bn_std - - def _denormalize_patch_latents(self, latents): - latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) - latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(latents.device, latents.dtype) - latents = latents * latents_bn_std + latents_bn_mean - return Flux2KleinPipeline._unpatchify_latents(latents) - - def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) - latent = self.vae.encode(image).latent_dist.sample() - return self._normalize_patch_latents(latent) - - def encode_condition(self, sample): - prompt = sample["conditioning"]["prompt"] - return self.encode_prompt_condition(prompt) - - def encode_prompt_condition(self, prompt): - model_config = self.config["model"] - prompt_embed, text_ids = self.text_pipeline.encode_prompt( - prompt=prompt, - device=self.device, - num_images_per_prompt=1, - max_sequence_length=model_config.get("max_sequence_length", 512), - text_encoder_out_layers=tuple(model_config.get("text_encoder_out_layers", (9, 18, 27))), - ) - return {"prompt_embed": prompt_embed, "text_ids": text_ids} - - def encode_prompt_text(self, prompt): - return self.encode_prompt_condition(prompt) - - def dmd_latent_shape(self, batch_size, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return ( - batch_size, - self.transformer.config.in_channels, - latent_h // 2, - latent_w // 2, - ) - - def prepare_denoiser_input(self, noisy_latent, condition=None): - h, w = noisy_latent.shape[2], noisy_latent.shape[3] - packed = Flux2KleinPipeline._pack_latents(noisy_latent) - img_ids = Flux2KleinPipeline._prepare_latent_ids(noisy_latent).to(self.device) - return Flux2KleinDenoiserInput( - hidden_states=packed, - img_ids=img_ids, - height=h, - width=w, - ) - - def denoise(self, denoiser_input, timestep_or_sigma, condition): - return self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, - guidance=None, - encoder_hidden_states=condition["prompt_embed"], - txt_ids=condition["text_ids"], - img_ids=denoiser_input.img_ids, - joint_attention_kwargs={}, - return_dict=False, - )[0] - - def postprocess_denoiser_output(self, prediction, denoiser_input): - return Flux2KleinPipeline._unpack_latents_with_ids( - prediction, - denoiser_input.img_ids, - height=denoiser_input.height, - width=denoiser_input.width, - ) - - def prepare_infer_latents(self, height, width, generator=None): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - shape = (1, self.transformer.config.in_channels, latent_h // 2, latent_w // 2) - return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) - - def decode_latent(self, latent): - latent = self._denormalize_patch_latents(latent) - image = self.vae.decode(latent).sample - return self.image_processor.postprocess(image, output_type="pil") - - @torch.no_grad() - def decode_packed_x0_to_images(self, packed_x0, latent_ids): - # height/width must be latent token grid sizes from img_ids, not pixel sizes. - unpatchified = Flux2KleinPipeline._unpack_latents_with_ids(packed_x0, latent_ids) - return self.decode_latent(unpatchified) - - def assemble_pipeline(self, scheduler=None): - return Flux2KleinPipeline( - tokenizer=self.text_pipeline.tokenizer, - text_encoder=self.text_pipeline.text_encoder, - vae=self.vae, - transformer=self.transformer, - scheduler=scheduler or self.text_pipeline.scheduler, - is_distilled=self.text_pipeline.config.is_distilled, - ).to(self.device) - - def get_pipeline_infer_kwargs(self, infer_config): - enable_cfg = infer_config.get("enable_cfg", True) - return { - "height": infer_config.get("height", infer_config.get("default_height", 1024)), - "width": infer_config.get("width", infer_config.get("default_width", 1024)), - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", 4.0) if enable_cfg else 1.0, - "max_sequence_length": self.config["model"].get("max_sequence_length", 512), - "text_encoder_out_layers": tuple(self.config["model"].get("text_encoder_out_layers", (9, 18, 27))), - } - - @torch.no_grad() - def prepare_reference_image_latents(self, images): - pipeline = self.assemble_pipeline() - reference_images = images.to(device=self.device, dtype=self.running_dtype) - encoded_image_latents = pipeline._encode_vae_image(image=reference_images, generator=None) - image_latent_ids = pipeline._prepare_image_ids([encoded_image_latents[:1]]) - image_latent_ids = image_latent_ids.repeat(encoded_image_latents.shape[0], 1, 1).to(self.device) - image_latents = pipeline._pack_latents(encoded_image_latents).to(device=self.device, dtype=self.running_dtype) - return image_latents, image_latent_ids - - def prepare_dopsd_initial_latents(self, height, width, batch_size, generator=None): - pipeline = self.assemble_pipeline() - num_latents_channels = self.transformer.config.in_channels // 4 - return pipeline.prepare_latents( - batch_size=batch_size, - num_latents_channels=num_latents_channels, - height=height, - width=width, - dtype=self.running_dtype, - device=self.device, - generator=generator, - latents=None, - ) - - def predict_velocity(self, packed_latents, timestep, condition, latent_ids, adapter_name, teacher_image_latents=None, teacher_image_latent_ids=None): - self.set_active_adapter(adapter_name) - if teacher_image_latents is not None: - hidden_states = torch.cat([packed_latents, teacher_image_latents], dim=1) - img_ids = torch.cat([latent_ids, teacher_image_latent_ids], dim=1) - else: - hidden_states = packed_latents - img_ids = latent_ids - - v_pred = self.transformer( - hidden_states=hidden_states, - timestep=timestep, - guidance=None, - encoder_hidden_states=condition["prompt_embed"], - txt_ids=condition["text_ids"], - img_ids=img_ids, - joint_attention_kwargs={}, - return_dict=False, - )[0] - return v_pred[:, : packed_latents.size(1)] diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image.py deleted file mode 100644 index 68cf5b990..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image.py +++ /dev/null @@ -1,190 +0,0 @@ -from dataclasses import dataclass - -import torch -from diffusers import AutoencoderKL, LongCatImagePipeline -from diffusers.image_processor import VaeImageProcessor -from diffusers.models.transformers import LongCatImageTransformer2DModel -from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids - -from lightx2v_train.utils.registry import MODEL_REGISTER - -from .base import BaseModel - - -@dataclass -class LongCatImageDenoiserInput: - hidden_states: torch.Tensor - img_ids: torch.Tensor - height: int - width: int - - -@MODEL_REGISTER("longcat_image") -class LongCatImageModel(BaseModel): - pipeline_cls = LongCatImagePipeline - - def load_components(self, transformer_only=False, reference_model=None): - if transformer_only: - if reference_model is not None: - self.text_pipeline = reference_model.text_pipeline - self.vae = reference_model.vae - self.image_processor = reference_model.image_processor - self.transformer = self.load_transformer() - self._maybe_set_attention_backend() - return - - model_path = self.config["model"]["pretrained_model_name_or_path"] - self.text_pipeline = LongCatImagePipeline.from_pretrained( - model_path, - transformer=None, - vae=None, - torch_dtype=self.running_dtype, - ).to(self.device) - self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(self.device, dtype=self.running_dtype) - self.transformer = self.load_transformer() - self.text_pipeline.text_encoder.requires_grad_(False) - self.vae.requires_grad_(False) - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - self._maybe_set_attention_backend() - - def load_transformer(self, model_path=None): - model_path = model_path or self.config["model"]["pretrained_model_name_or_path"] - return LongCatImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) - - def _maybe_set_attention_backend(self): - attention_backend = self.config["model"].get("attention_backend", None) - if attention_backend is not None: - self.transformer.set_attention_backend(attention_backend) - - def denoiser_module(self): - return self.transformer - - def fsdp2_shard_plan(self, fsdp_config): - reshard_config = fsdp_config["reshard_after_forward"] - return [ - { - "modules": self.transformer.transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "modules": self.transformer.single_transformer_blocks, - "reshard_after_forward": reshard_config["block_reshard"], - }, - { - "module": self.transformer, - "reshard_after_forward": reshard_config["root_reshard"], - }, - ] - - @property - def vae_scale_factor(self): - return 2 ** (len(self.vae.config.block_out_channels) - 1) - - def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) - latent = self.vae.encode(image).latent_dist.sample() - shift = getattr(self.vae.config, "shift_factor", 0.0) - scale = getattr(self.vae.config, "scaling_factor", 1.0) - return (latent - shift) * scale - - def encode_condition(self, sample): - prompt = sample["conditioning"]["prompt"] - return self.encode_prompt_condition(prompt) - - def encode_prompt_condition(self, prompt): - rewrite_training = self.config.get( - "enable_prompt_rewrite_training", - self.config["model"].get("enable_prompt_rewrite_training", False), - ) - if rewrite_training: - prompt = self.text_pipeline.rewire_prompt(prompt, self.device) - prompt_embed, text_ids = self.text_pipeline.encode_prompt( - prompt=prompt, - num_images_per_prompt=1, - ) - return {"prompt_embed": prompt_embed, "text_ids": text_ids} - - def prepare_denoiser_input(self, noisy_latent, condition=None): - n = noisy_latent.shape[0] - h, w = noisy_latent.shape[2], noisy_latent.shape[3] - packed = LongCatImagePipeline._pack_latents(noisy_latent, n, noisy_latent.shape[1], h, w) - img_ids = prepare_pos_ids( - modality_id=1, - type="image", - start=(self.text_pipeline.tokenizer_max_length, self.text_pipeline.tokenizer_max_length), - height=h // 2, - width=w // 2, - ).to(self.device) - return LongCatImageDenoiserInput( - hidden_states=packed, - img_ids=img_ids, - height=h, - width=w, - ) - - def denoise(self, denoiser_input, timestep_or_sigma, condition): - return self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, - guidance=None, - encoder_hidden_states=condition["prompt_embed"], - txt_ids=condition["text_ids"], - img_ids=denoiser_input.img_ids, - return_dict=False, - )[0] - - def postprocess_denoiser_output(self, prediction, denoiser_input): - return LongCatImagePipeline._unpack_latents( - prediction, - height=denoiser_input.height * self.vae_scale_factor, - width=denoiser_input.width * self.vae_scale_factor, - vae_scale_factor=self.vae_scale_factor, - ) - - def prepare_infer_latents(self, height, width, generator=None): - latent_h = height // self.vae_scale_factor - latent_w = width // self.vae_scale_factor - # latent shape: (batch=1, latent_channels, latent_h, latent_w) - shape = (1, self.vae.config.latent_channels, latent_h, latent_w) - return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) - - def dmd_latent_shape(self, batch_size, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return (int(batch_size), int(self.vae.config.latent_channels), latent_h, latent_w) - - def cfg_on_denoiser_output(self): - return True - - def decode_latent(self, latent): - # Reverse the normalization from encode_to_latent: - # encode: normalized = (raw - shift) * scale - # decode: raw = normalized / scale + shift - shift = getattr(self.vae.config, "shift_factor", 0.0) - scale = getattr(self.vae.config, "scaling_factor", 1.0) - latent = latent / scale + shift - - image = self.vae.decode(latent).sample # (B, C, H, W) - return self.image_processor.postprocess(image, output_type="pil") - - def assemble_pipeline(self, scheduler=None): - return LongCatImagePipeline( - tokenizer=self.text_pipeline.tokenizer, - text_encoder=self.text_pipeline.text_encoder, - text_processor=self.text_pipeline.text_processor, - vae=self.vae, - transformer=self.transformer, - scheduler=scheduler or self.text_pipeline.scheduler, - ).to(self.device) - - def get_pipeline_infer_kwargs(self, infer_config): - enable_cfg = infer_config.get("enable_cfg", False) - return { - "height": infer_config.get("height", infer_config.get("default_height", 1024)), - "width": infer_config.get("width", infer_config.get("default_width", 1024)), - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", 4.0) if enable_cfg else 1.0, - "enable_cfg_renorm": infer_config.get("enable_cfg_renorm", True), - "cfg_renorm_min": infer_config.get("cfg_renorm_min", 0.0), - "enable_prompt_rewrite": infer_config.get("enable_prompt_rewrite", True), - } diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/__init__.py new file mode 100644 index 000000000..2f8417c49 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/__init__.py @@ -0,0 +1 @@ +"""LongCat Image model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/__init__.py new file mode 100644 index 000000000..a62dbc0aa --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/__init__.py @@ -0,0 +1,6 @@ +"""Capabilities implemented specifically by LongCat-Image.""" + +from .longcat_image_consistency_model_capability import LongCatImageConsistencyModelCapability +from .longcat_image_edit_distribution_matching_capability import LongCatImageEditDistributionMatchingCapability + +__all__ = ["LongCatImageConsistencyModelCapability", "LongCatImageEditDistributionMatchingCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py new file mode 100644 index 000000000..2ea350ef5 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py @@ -0,0 +1,22 @@ +"""Consistency-model capability for LongCat-Image.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( + ProjectedTimeEmbeddingAdapter, + TimeConditionedConsistencyModelCapability, +) + + +class LongCatImageConsistencyModelCapability(TimeConditionedConsistencyModelCapability): + """Bind generic consistency extensions to LongCat's time embedding.""" + + def __init__(self, model) -> None: + super().__init__( + model, + ProjectedTimeEmbeddingAdapter( + hook_module_path="time_embed", + projection_module_path="time_embed.time_proj", + embedding_module_path="time_embed.timestep_embedder", + embedding_dimension_path="inner_dim", + time_scale=1000.0, + ), + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_edit_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_edit_distribution_matching_capability.py new file mode 100644 index 000000000..2147e57c7 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_edit_distribution_matching_capability.py @@ -0,0 +1,30 @@ +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _cached_condition_pair, + _negative_prompt, + _require_single_prompt, +) + + +class LongCatImageEditDistributionMatchingCapability(GenericDistributionMatchingCapability): + cache_uses_sample_context = True + + def encode_conditions(self, batch, negative_prompt, guidance_scale, broadcast): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + cached = _cached_condition_pair(batch, self.model, require_negative=guidance_scale > 1) + if cached is not None: + positive, negative = cached + return broadcast(positive), broadcast(negative) if negative is not None else None + prompts = [prompt] + if guidance_scale > 1: + prompts.append(_negative_prompt(conditioning, negative_prompt, scalar=scalar)) + + with torch.no_grad(): + conditions = self.model.encode_conditions_with_source(batch, prompts) + positive = conditions[0] + negative = conditions[1] if len(conditions) > 1 else None + return broadcast(positive), broadcast(negative) if negative is not None else None diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py new file mode 100644 index 000000000..5ac13b1da --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py @@ -0,0 +1,163 @@ +import math + +import torch +from diffusers import AutoencoderKL +from diffusers.image_processor import VaeImageProcessor + +from lightx2v_train.utils.image_ops import ( + align_dimension, + calculate_area_dimensions, + pil_to_image_tensor, + resize_and_center_crop, +) +from lightx2v_train.utils.registry import SAMPLE_PROCESSOR_REGISTER + + +def _size_multiple_from_config(config): + processor_config = config.get("data", {}).get("processor", {}) + preprocessing = config.get("model", {}).get("input_preprocessing", {}) + value = processor_config.get("size_multiple", preprocessing.get("size_multiple")) + if value is None: + model_path = config["model"]["pretrained_model_name_or_path"] + vae_config = AutoencoderKL.load_config(model_path, subfolder="vae") + value = 2 ** (len(vae_config["block_out_channels"]) - 1) * 2 + value = int(value) + if value <= 0: + raise ValueError(f"size_multiple must be positive, got {value}") + return value + + +def _target_area_from_config(config): + processor_config = config.get("data", {}).get("processor", {}) + if "target_area" in processor_config: + return int(processor_config["target_area"]) + preprocessing = config.get("model", {}).get("input_preprocessing", {}) + if "target_area" in preprocessing: + return int(preprocessing["target_area"]) + data_config = config.get("data", {}) + for split in ("train", "val"): + if "target_area" in data_config.get(split, {}): + return int(data_config[split]["target_area"]) + return 1024 * 1024 + + +def _optional_scalar(mapping, key): + value = mapping.get(key) + if value is None: + return None + if torch.is_tensor(value): + if value.numel() != 1: + raise ValueError(f"{key} must contain one value, got {value.numel()}") + value = value.item() + elif isinstance(value, (list, tuple)): + if len(value) != 1: + raise ValueError(f"{key} must contain one value, got {len(value)}") + value = value[0] + if torch.is_tensor(value): + value = value.item() + return int(value) + + +def _explicit_target_size(sample): + meta = sample.get("meta", {}) + height = _optional_scalar(meta, "target_height") + width = _optional_scalar(meta, "target_width") + if (height is None) != (width is None): + raise ValueError("meta.target_height and meta.target_width must be provided together") + return height, width + + +def _align_up(value, multiple): + if value <= 0: + raise ValueError(f"Image dimensions must be positive, got {value}") + return math.ceil(value / multiple) * multiple + + +def _calculate_edit_dimensions(target_area, ratio, multiple): + if target_area <= 0: + raise ValueError(f"target_area must be positive, got {target_area}") + if ratio <= 0: + raise ValueError(f"Image aspect ratio must be positive, got {ratio}") + width = _align_up(math.sqrt(target_area * ratio), multiple) + height = _align_up(math.sqrt(target_area / ratio), multiple) + return width, height + + +@SAMPLE_PROCESSOR_REGISTER("longcat_image") +class LongCatImageDataProcessor: + def __init__(self, config): + self.image_processor = VaeImageProcessor(vae_scale_factor=_size_multiple_from_config(config)) + self.target_area = _target_area_from_config(config) + self.unconditional_prompt = config["model"].get("unconditional_prompt", " ") + + def __call__(self, sample): + inputs = sample["inputs"] + image = inputs.pop("target_image", None) + if image is None: + return sample + + height, width = _explicit_target_size(sample) + multiple = int(self.image_processor.config.vae_scale_factor) + if height is None: + ratio = image.width / image.height + width, height = calculate_area_dimensions(self.target_area, ratio, multiple) + else: + width = align_dimension(width, multiple) + height = align_dimension(height, multiple) + image = resize_and_center_crop(image, width, height) + inputs["target_pixel_values"] = self.image_processor.preprocess(image)[0] + return sample + + def infer_target_size(self, sample, default_height, default_width): + height, width = _explicit_target_size(sample) + height = default_height if height is None else height + width = default_width if width is None else width + multiple = int(self.image_processor.config.vae_scale_factor) + return align_dimension(int(height), multiple), align_dimension(int(width), multiple) + + +@SAMPLE_PROCESSOR_REGISTER("longcat_image_edit") +class LongCatImageEditDataProcessor(LongCatImageDataProcessor): + def __init__(self, config): + super().__init__(config) + self.unconditional_prompt = config["model"].get("unconditional_prompt", "") + + def __call__(self, sample): + inputs = sample["inputs"] + target_image = inputs.pop("target_image", None) + source_images = inputs.pop("source_images", []) + if len(source_images) != 1: + raise ValueError(f"LongCat Image Edit requires exactly one source image, got {len(source_images)}") + + source_image = source_images[0] + height, width = self._resolve_edit_size(sample, source_image) + sample["meta"]["target_height"] = height + sample["meta"]["target_width"] = width + sample["meta"]["reference_image_height"] = source_image.height + sample["meta"]["reference_image_width"] = source_image.width + + source_image = self.image_processor.resize(source_image, height, width) + prompt_image = self.image_processor.resize(source_image, height // 2, width // 2) + inputs["source_condition_image"] = pil_to_image_tensor(prompt_image) + inputs["source_vae_pixel_values"] = self.image_processor.preprocess(source_image)[0] + + if target_image is not None: + target_image = resize_and_center_crop(target_image, width, height) + inputs["target_pixel_values"] = self.image_processor.preprocess(target_image)[0] + return sample + + def infer_target_size(self, sample, default_height, default_width): + del default_height, default_width + height, width = _explicit_target_size(sample) + if height is None: + raise ValueError("LongCat Image Edit requires source-derived target dimensions") + return height, width + + def _resolve_edit_size(self, sample, source_image): + height, width = _explicit_target_size(sample) + multiple = int(self.image_processor.config.vae_scale_factor) + if height is not None: + return _align_up(height, multiple), _align_up(width, multiple) + ratio = source_image.width / source_image.height + width, height = _calculate_edit_dimensions(self.target_area, ratio, multiple) + return height, width diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py new file mode 100644 index 000000000..4c7ed5907 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -0,0 +1,269 @@ +from dataclasses import dataclass + +import torch +from diffusers import AutoencoderKL, LongCatImagePipeline +from diffusers.image_processor import VaeImageProcessor +from diffusers.models.transformers import LongCatImageTransformer2DModel +from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids + +from lightx2v_train.model_capabilities import ConsistencyModelCapability, DistributionMatchingCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.longcat_image.capability_adapters import LongCatImageConsistencyModelCapability +from lightx2v_train.utils.registry import MODEL_REGISTER + +from ..base import BaseModel + + +@dataclass +class LongCatImageDenoiserInput: + hidden_states: torch.Tensor + img_ids: torch.Tensor + target_token_length: int + height: int + width: int + + +@MODEL_REGISTER("longcat_image") +class LongCatImageModel(BaseModel): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/meituan-longcat/LongCat-Image + - https://huggingface.co/meituan-longcat/LongCat-Image-Dev + """ + + pipeline_cls = LongCatImagePipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability + supports_cfg_renorm = True + supports_prompt_rewrite = True + + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistributionMatchingCapability, + self.distribution_matching_capability_cls( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="latent_channels", + ), + guidance_in_denoiser_space=True, + ), + ) + self.capabilities.register( + ConsistencyModelCapability, + LongCatImageConsistencyModelCapability(self), + ) + + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + model_path = self.config["model"]["pretrained_model_name_or_path"] + if load_condition_encoder: + self._load_condition_encoder(model_path) + if load_vae: + self._load_vae(model_path) + else: + self._load_vae_config(model_path) + if load_transformer: + self.transformer = self.load_transformer() + self._maybe_set_attention_backend() + + def _load_condition_encoder(self, model_path): + self.text_pipeline = self.pipeline_cls.from_pretrained( + model_path, + transformer=None, + vae=None, + torch_dtype=self.running_dtype, + ).to(self.device) + self.text_pipeline.text_encoder.requires_grad_(False) + self.text_pipeline.text_encoder.eval() + + def _load_vae(self, model_path): + self.vae = AutoencoderKL.from_pretrained( + model_path, + subfolder="vae", + torch_dtype=self.running_dtype, + ).to(self.device) + self.vae_config = self.vae.config + self.vae.requires_grad_(False) + self.vae.eval() + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + + def _load_vae_config(self, model_path): + self.vae_config = AutoencoderKL.load_config(model_path, subfolder="vae") + + @property + def latent_channels(self): + return int(self.vae_config["latent_channels"] if isinstance(self.vae_config, dict) else self.vae_config.latent_channels) + + def load_transformer(self, model_path=None): + model_path = model_path or self.config["model"]["pretrained_model_name_or_path"] + return LongCatImageTransformer2DModel.from_pretrained( + model_path, + subfolder="transformer", + torch_dtype=self.running_dtype, + ).to(self.device) + + def _maybe_set_attention_backend(self): + attention_backend = self.config["model"].get("attention_backend") + if attention_backend is not None: + self.transformer.set_attention_backend(attention_backend) + + def denoiser_module(self): + return self.transformer + + def fsdp2_shard_plan(self, fsdp_config): + reshard_config = fsdp_config["reshard_after_forward"] + return [ + { + "modules": self.transformer.transformer_blocks, + "reshard_after_forward": reshard_config["block_reshard"], + }, + { + "modules": self.transformer.single_transformer_blocks, + "reshard_after_forward": reshard_config["block_reshard"], + }, + { + "module": self.transformer, + "reshard_after_forward": reshard_config["root_reshard"], + }, + ] + + @property + def vae_scale_factor(self): + config = self.vae_config + block_out_channels = config["block_out_channels"] if isinstance(config, dict) else config.block_out_channels + return 2 ** (len(block_out_channels) - 1) + + def _normalize_latents(self, latents): + shift = getattr(self.vae.config, "shift_factor", 0.0) + scale = getattr(self.vae.config, "scaling_factor", 1.0) + return (latents - shift) * scale + + def _denormalize_latents(self, latents): + shift = getattr(self.vae.config, "shift_factor", 0.0) + scale = getattr(self.vae.config, "scaling_factor", 1.0) + return latents / scale + shift + + def encode_to_latent(self, sample): + return self._encode_target_latent(sample, mode="sample") + + def encode_to_cache_latent(self, sample): + return self._encode_target_latent(sample, mode="mode") + + def _encode_target_latent(self, sample, *, mode): + image = sample["inputs"]["target_pixel_values"] + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4: + raise ValueError(f"Expected target_pixel_values with shape [B, C, H, W], got {tuple(image.shape)}") + image = image.to(device=self.device, dtype=self.running_dtype) + latent = getattr(self.vae.encode(image).latent_dist, mode)() + return self._normalize_latents(latent) + + def encode_condition(self, sample): + prompt = sample["conditioning"]["prompt"] + if self.config["model"].get("enable_prompt_rewrite_training", False): + prompt = self.text_pipeline.rewire_prompt(prompt, self.device) + return self.encode_prompt_condition(prompt) + + def encode_inference_condition(self, sample, *, is_negative=False): + prompt = sample["conditioning"]["prompt"] + rewrite = self.config.get("inference", {}).get("enable_prompt_rewrite", True) + if self.supports_prompt_rewrite and rewrite and not is_negative: + prompt = self.text_pipeline.rewire_prompt(prompt, self.device) + return self.encode_prompt_condition(prompt) + + def encode_prompt_condition(self, prompt): + prompt_embed, text_ids = self.text_pipeline.encode_prompt( + prompt=prompt, + num_images_per_prompt=1, + ) + return {"prompt_embed": prompt_embed, "text_ids": text_ids} + + def _get_additional_image_condition(self, condition, position_start): + del condition, position_start + return None, None + + def prepare_denoiser_input(self, noisy_latent, condition=None): + if condition is None: + raise ValueError(f"{type(self).__name__}.prepare_denoiser_input requires condition") + + batch_size, channels, height, width = noisy_latent.shape + target_tokens = self.pipeline_cls._pack_latents(noisy_latent, batch_size, channels, height, width) + position_start = int(condition["prompt_embed"].shape[1]) + target_ids = prepare_pos_ids( + modality_id=1, + type="image", + start=(position_start, position_start), + height=height // 2, + width=width // 2, + ).to(self.device) + + additional_tokens, additional_ids = self._get_additional_image_condition(condition, position_start) + hidden_states = target_tokens + img_ids = target_ids + if additional_tokens is not None: + if additional_ids is None: + raise ValueError("Additional LongCat image tokens require matching image IDs") + if additional_tokens.shape[0] != batch_size or additional_tokens.shape[2] != target_tokens.shape[2]: + raise ValueError(f"Additional LongCat image tokens must match the target batch and channel dimensions, got {tuple(additional_tokens.shape)} and {tuple(target_tokens.shape)}") + hidden_states = torch.cat([target_tokens, additional_tokens], dim=1) + img_ids = torch.cat([target_ids, additional_ids], dim=0) + + return LongCatImageDenoiserInput( + hidden_states=hidden_states, + img_ids=img_ids, + target_token_length=target_tokens.shape[1], + height=height, + width=width, + ) + + def denoise(self, denoiser_input, timestep_or_sigma, condition): + prediction = self.transformer( + hidden_states=denoiser_input.hidden_states, + timestep=timestep_or_sigma, + guidance=None, + encoder_hidden_states=condition["prompt_embed"], + txt_ids=condition["text_ids"], + img_ids=denoiser_input.img_ids, + return_dict=False, + )[0] + return prediction[:, : denoiser_input.target_token_length] + + def postprocess_denoiser_output(self, prediction, denoiser_input): + return self.pipeline_cls._unpack_latents( + prediction, + height=denoiser_input.height * self.vae_scale_factor, + width=denoiser_input.width * self.vae_scale_factor, + vae_scale_factor=self.vae_scale_factor, + ) + + def apply_cfg(self, positive, negative, guidance_scale): + prediction = super().apply_cfg(positive, negative, guidance_scale) + infer_config = self.config.get("inference", {}) + if not self.supports_cfg_renorm or not infer_config.get("enable_cfg_renorm", True): + return prediction + + positive_norm = torch.norm(positive, dim=-1, keepdim=True) + prediction_norm = torch.norm(prediction, dim=-1, keepdim=True) + minimum = float(infer_config.get("cfg_renorm_min", 0.0)) + scale = (positive_norm / prediction_norm.clamp_min(1e-8)).clamp(min=minimum, max=1.0) + return prediction * scale + + def prepare_infer_latents(self, height, width, generator=None): + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + shape = (1, self.latent_channels, latent_height, latent_width) + return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) + + def decode_latent(self, latent): + image = self.vae.decode(self._denormalize_latents(latent)).sample + return self.image_processor.postprocess(image, output_type="pil") diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py new file mode 100644 index 000000000..e2eabdb78 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py @@ -0,0 +1,85 @@ +from diffusers import LongCatImageEditPipeline +from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids + +from lightx2v_train.model_zoo.longcat_image.capability_adapters import LongCatImageEditDistributionMatchingCapability +from lightx2v_train.utils.image_ops import image_tensor_to_pil +from lightx2v_train.utils.registry import MODEL_REGISTER + +from .longcat_image import LongCatImageModel + + +@MODEL_REGISTER("longcat_image_edit") +class LongCatImageEditModel(LongCatImageModel): + """Supports weights from these Hugging Face repos: + - https://huggingface.co/meituan-longcat/LongCat-Image-Edit + """ + + pipeline_cls = LongCatImageEditPipeline + distribution_matching_capability_cls = LongCatImageEditDistributionMatchingCapability + supports_cfg_renorm = False + supports_prompt_rewrite = False + default_unconditional_prompt = "" + shared_condition_keys = ("source_tokens", "source_height", "source_width") + + def encode_condition(self, sample): + prompt = sample["conditioning"]["prompt"] + return self.encode_conditions_with_source(sample, [prompt])[0] + + def encode_conditions_with_context(self, sample, prompts): + return self.encode_conditions_with_source(sample, prompts) + + def encode_inference_condition(self, sample, *, is_negative=False): + del is_negative + return self.encode_condition(sample) + + def encode_conditions_with_source(self, sample, prompts): + inputs = sample["inputs"] + source_image = inputs.get("source_condition_image") + source_pixels = inputs.get("source_vae_pixel_values") + if source_image is None or source_pixels is None: + raise ValueError("LongCat Image Edit requires source_condition_image and source_vae_pixel_values") + + source_image = image_tensor_to_pil(source_image) + source_tokens, source_height, source_width = self._encode_source_latents(source_pixels) + conditions = [] + for prompt in prompts: + prompt_embed, text_ids = self.text_pipeline.encode_prompt( + prompt=prompt, + image=source_image, + num_images_per_prompt=1, + ) + conditions.append( + { + "prompt_embed": prompt_embed, + "text_ids": text_ids, + "source_tokens": source_tokens, + "source_height": source_height, + "source_width": source_width, + } + ) + return conditions + + def _encode_source_latents(self, source_pixels): + if source_pixels.ndim == 3: + source_pixels = source_pixels.unsqueeze(0) + if source_pixels.ndim != 4 or source_pixels.shape[0] != 1: + raise ValueError(f"Expected one source image with shape [1, C, H, W], got {tuple(source_pixels.shape)}") + + source_pixels = source_pixels.to(device=self.device, dtype=self.running_dtype) + latents = self._normalize_latents(self.vae.encode(source_pixels).latent_dist.mode()) + batch_size, channels, height, width = latents.shape + tokens = self.pipeline_cls._pack_latents(latents, batch_size, channels, height, width) + return tokens, height, width + + def _get_additional_image_condition(self, condition, position_start): + source_tokens = condition.get("source_tokens") + if source_tokens is None: + raise ValueError("LongCat Image Edit condition is missing source_tokens") + source_ids = prepare_pos_ids( + modality_id=2, + type="image", + start=(position_start, position_start), + height=int(condition["source_height"]) // 2, + width=int(condition["source_width"]) // 2, + ).to(self.device) + return source_tokens, source_ids diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py deleted file mode 100644 index 1fcc3061e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py +++ /dev/null @@ -1,371 +0,0 @@ -import json -import os -from contextlib import nullcontext - -import torch -from loguru import logger -from peft import LoraConfig, inject_adapter_in_model -from peft.utils import set_peft_model_state_dict -from safetensors import safe_open -from safetensors.torch import load_file - -from lightx2v_train.model_zoo.native.ltx2 import ( - EMBEDDINGS_PROCESSOR_KEY_OPS, - GEMMA_LLM_KEY_OPS, - GEMMA_MODEL_OPS, - LTXV_MODEL_COMFY_RENAMING_MAP, - EmbeddingsProcessorConfigurator, - GemmaTextEncoderConfigurator, - LTXModelConfigurator, - SingleGPUModelBuilder, - convert_to_additive_mask, - find_matching_file, - module_ops_from_gemma_root, -) -from lightx2v_train.utils.registry import MODEL_REGISTER -from lightx2v_train.utils.utils import get_running_dtype - -from .base import BaseModel - - -@MODEL_REGISTER("ltx_t2av_ar") -@MODEL_REGISTER("ltx_t2av") -class LTX2T2AVModel(BaseModel): - pipeline_cls = None - - def load_components(self, transformer_only=False, reference_model=None): - model_config = self.config["model"] - model_path = model_config["pretrained_model_name_or_path"] - self.pretrained_model_path = os.path.abspath(os.path.expanduser(str(model_path))) - self._fsdp2_activation_checkpointing = False - - self.transformer_param_dtype = get_running_dtype(model_config.get("transformer_param_dtype", "bf16")) - self.embeddings_dtype = get_running_dtype(model_config.get("embeddings_dtype", "bf16")) - self.text_encoder_dtype = get_running_dtype(model_config.get("text_encoder_dtype", "bf16")) - self.load_transformer = bool(model_config.get("load_transformer", True)) - self.load_embeddings_processor = bool(model_config.get("load_embeddings_processor", True)) - self.load_text_encoder = bool(model_config.get("load_text_encoder", False)) - self.text_encoder_cpu = bool(model_config.get("text_encoder_cpu", False)) - self.load_text_encoder_in_8bit = bool(model_config.get("load_text_encoder_in_8bit", False)) - self.text_encoder_path = model_config.get("text_encoder_path", "/data/nvme0/gushiqiao/models/official_models/LTX-2") - training_config = self.config.get("training", {}) - teacher_forcing_config = training_config.get("teacher_forcing", {}) - self.num_frame_per_chunk = int(model_config.get("num_frame_per_chunk", teacher_forcing_config.get("num_frame_per_chunk", 1))) - self.local_attn_size = int(model_config.get("local_attn_size", training_config.get("dmd", {}).get("local_attn_size", -1))) - self.sink_size = int(model_config.get("sink_size", training_config.get("dmd", {}).get("sink_size", 0))) - if "defer_kv_cache_updates" in model_config: - self.defer_kv_cache_updates = bool(model_config["defer_kv_cache_updates"]) - else: - self.defer_kv_cache_updates = bool(model_config.get("defer_cache_updates", False)) - self.detach_kv_cache_updates = bool(model_config.get("detach_kv_cache_updates", False)) - - if transformer_only: - self.load_embeddings_processor = False - self.load_text_encoder = False - if reference_model is not None: - self.embeddings_processor = reference_model.embeddings_processor - self.text_encoder = getattr(reference_model, "text_encoder", None) - self.text_encoder_path = getattr(reference_model, "text_encoder_path", self.text_encoder_path) - self.text_encoder_cpu = getattr(reference_model, "text_encoder_cpu", self.text_encoder_cpu) - - transformer_path = self._resolve_transformer_path(model_path, model_config, transformer_only) - if self.load_transformer: - self.transformer = self._load_transformer(transformer_path) - self._configure_causal_transformer() - else: - self.transformer = None - - self.embeddings_processor = getattr(self, "embeddings_processor", None) - self.text_encoder = getattr(self, "text_encoder", None) - if self.load_embeddings_processor: - self.embeddings_processor = self._load_embeddings_processor(model_path) - self.embeddings_processor.requires_grad_(False) - self.embeddings_processor.eval() - if hasattr(self.embeddings_processor, "feature_extractor") and not self.load_text_encoder: - self.embeddings_processor.feature_extractor = None - if self.load_text_encoder: - self.text_encoder = self._load_text_encoder(self.text_encoder_path) - self.text_encoder.requires_grad_(False) - self.text_encoder.eval() - - def _load_transformer(self, model_path): - builder = SingleGPUModelBuilder( - model_path=str(model_path), - model_class_configurator=LTXModelConfigurator, - model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP, - ) - self.transformer_checkpoint_config = builder.model_config() - transformer = builder.build(device=self.device, dtype=self.transformer_param_dtype) - return transformer.to(self.device, dtype=self.transformer_param_dtype) - - def _resolve_transformer_path(self, model_path, model_config, transformer_only): - self.loaded_student_checkpoint_path = None - if transformer_only or model_config.get("name") != "ltx_t2av_ar": - return model_path - - student_config = model_config.get("student", {}) - checkpoint_path = student_config.get("checkpoint_path") if isinstance(student_config, dict) else None - if checkpoint_path is None: - return model_path - - checkpoint_path = os.path.abspath(os.path.expanduser(str(checkpoint_path))) - candidates = [checkpoint_path] - if os.path.isdir(checkpoint_path): - candidates = [ - os.path.join(checkpoint_path, "consolidated.safetensors"), - os.path.join(checkpoint_path, "student_consolidated.safetensors"), - ] - for candidate in candidates: - if os.path.isfile(candidate) and candidate.endswith(".safetensors"): - self.loaded_student_checkpoint_path = checkpoint_path - logger.info("[train] loading LTX2 AR student transformer from {}", candidate) - return candidate - return model_path - - def prepare_consolidated_state_dict(self, state_dict): - consolidated = {} - for key, value in state_dict.items(): - while key.startswith("module.") or key.startswith("_fsdp_wrapped_module."): - key = key.split(".", 1)[1] - key = key.replace("._checkpoint_wrapped_module.", ".") - if not key.startswith("model.diffusion_model."): - key = f"model.diffusion_model.{key}" - if key in consolidated: - raise ValueError(f"Duplicate LTX2 consolidated checkpoint key: {key}") - consolidated[key] = value - - if not os.path.isfile(self.pretrained_model_path): - raise FileNotFoundError(f"Official LTX checkpoint not found: {self.pretrained_model_path}") - added = 0 - with safe_open(self.pretrained_model_path, framework="pt", device="cpu") as handle: - for key in handle.keys(): - if key in consolidated: - continue - consolidated[key] = handle.get_tensor(key).contiguous() - added += 1 - logger.info("[checkpoint] added {} non-trained component tensors from {}", added, self.pretrained_model_path) - return consolidated - - def consolidated_safetensors_metadata(self): - metadata = super().consolidated_safetensors_metadata() - checkpoint_config = getattr(self, "transformer_checkpoint_config", None) - if checkpoint_config: - metadata["config"] = json.dumps(checkpoint_config, ensure_ascii=False, separators=(",", ":")) - return metadata - - def _configure_causal_transformer(self): - self.transformer.num_frame_per_block = self.num_frame_per_chunk - self.transformer.local_attn_size = self.local_attn_size - self.transformer.sink_size = self.sink_size - self.transformer.defer_kv_cache_updates = self.defer_kv_cache_updates - self.transformer.detach_kv_cache_updates = self.detach_kv_cache_updates - - def _load_embeddings_processor(self, model_path): - processor = SingleGPUModelBuilder( - model_path=str(model_path), - model_class_configurator=EmbeddingsProcessorConfigurator, - model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS, - ).build(device=self.device, dtype=self.embeddings_dtype) - return processor.to(self.device, dtype=self.embeddings_dtype) - - def _load_text_encoder(self, text_encoder_path): - if self.load_text_encoder_in_8bit: - try: - from ltx_trainer.gemma_8bit import load_8bit_gemma - except ImportError as exc: - raise ImportError("model.load_text_encoder_in_8bit=true requires ltx_trainer.gemma_8bit on PYTHONPATH. Disable it or add the LTX-2 trainer package.") from exc - - device = torch.device("cpu") if self.text_encoder_cpu else self.device - return load_8bit_gemma(text_encoder_path, self.text_encoder_dtype, device=device) - - gemma_model_folder = find_matching_file(str(text_encoder_path), "model*.safetensors").parent - gemma_weight_paths = tuple(str(path) for path in gemma_model_folder.rglob("*.safetensors")) - device = torch.device("cpu") if self.text_encoder_cpu else self.device - return SingleGPUModelBuilder( - model_path=gemma_weight_paths, - model_class_configurator=GemmaTextEncoderConfigurator, - model_sd_ops=GEMMA_LLM_KEY_OPS, - module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(text_encoder_path))), - ).build(device=device, dtype=self.text_encoder_dtype) - - def denoiser_module(self): - if self.transformer is None: - raise RuntimeError("LTX2 transformer is not loaded. Set model.load_transformer=True for training.") - return self.transformer - - def transformer_forward_context(self): - if self.device.type == "cuda" and self.running_dtype in {torch.float16, torch.bfloat16}: - return torch.autocast(device_type="cuda", dtype=self.running_dtype) - return nullcontext() - - def add_lora(self, rank, alpha, target_modules): - if not target_modules: - raise ValueError("training.lora.target_modules must be set for LTX2 LoRA training.") - lora_config = LoraConfig( - r=rank, - lora_alpha=alpha, - init_lora_weights="gaussian", - target_modules=target_modules, - ) - self._inject_lora(lora_config) - - def _inject_lora(self, lora_config, adapter_name="default"): - try: - self.transformer = inject_adapter_in_model(lora_config, self.transformer, adapter_name=adapter_name) - except TypeError: - self.transformer = inject_adapter_in_model(lora_config, self.transformer) - - def load_lora_for_infer(self, lora_path, adapter_name=None): - if adapter_name is None: - adapter_name = "default" - if not hasattr(self.denoiser_module(), "peft_config"): - training_config = self.config.get("training", {}) - lora_config = training_config.get("lora", {}) - rank = int(lora_config.get("rank", 64)) - alpha = int(lora_config.get("alpha", rank)) - target_modules = lora_config.get("target_modules") - self._inject_lora( - LoraConfig( - r=rank, - lora_alpha=alpha, - init_lora_weights="gaussian", - target_modules=target_modules, - ), - adapter_name=adapter_name, - ) - - weight_path = lora_path - if os.path.isdir(weight_path): - weight_path = os.path.join(weight_path, "pytorch_lora_weights.safetensors") - raw = load_file(weight_path) - peft_state_dict = {} - for key, value in raw.items(): - new_key = key.removeprefix("transformer.") - new_key = new_key.replace(".lora.down.weight", ".lora_A.weight") - new_key = new_key.replace(".lora.up.weight", ".lora_B.weight") - peft_state_dict[new_key] = value - - incompatible = set_peft_model_state_dict(self.denoiser_module(), peft_state_dict) - if incompatible and incompatible.unexpected_keys: - logger.warning("Unexpected keys when loading LTX2 LoRA: {}", incompatible.unexpected_keys) - self._infer_lora_adapter_name = adapter_name - - def enable_gradient_checkpointing(self): - if self._fsdp2_activation_checkpointing: - return - if hasattr(self.denoiser_module(), "set_gradient_checkpointing"): - self.denoiser_module().set_gradient_checkpointing(True) - - def _wrap_transformer_blocks_for_fsdp_checkpointing(self): - if self._fsdp2_activation_checkpointing: - return - - from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl, checkpoint_wrapper - - blocks = self.transformer.transformer_blocks - for block_idx, block in enumerate(blocks): - blocks[block_idx] = checkpoint_wrapper( - block, - checkpoint_impl=CheckpointImpl.NO_REENTRANT, - ) - self._fsdp2_activation_checkpointing = True - logger.info("Enabled FSDP2 activation checkpoint wrappers for {} LTX2 transformer blocks", len(blocks)) - - def fsdp2_shard_plan(self, fsdp_config): - training_config = self.config.get("training", {}) - if training_config.get("gradient_checkpointing", True): - self._wrap_transformer_blocks_for_fsdp_checkpointing() - reshard_config = fsdp_config.get( - "reshard_after_forward", - { - "root_reshard": False, - "block_reshard": True, - }, - ) - return [ - { - "modules": self.transformer.transformer_blocks, - "reshard_after_forward": reshard_config.get("block_reshard", True), - }, - { - "module": self.transformer, - "reshard_after_forward": reshard_config.get("root_reshard", False), - }, - ] - - def prepare_text_condition(self, conditions): - if self.embeddings_processor is None: - raise RuntimeError("LTX2 embeddings processor is not loaded. Set model.load_embeddings_processor=True.") - - if "video_prompt_embeds" in conditions: - video_features = conditions["video_prompt_embeds"] - audio_features = conditions.get("audio_prompt_embeds", video_features) - else: - video_features = conditions["prompt_embeds"] - audio_features = conditions["prompt_embeds"] - mask = conditions["prompt_attention_mask"] - video_features = self._ensure_prompt_feature_batch(video_features, "video prompt embeds").to(device=self.device) - if audio_features is not None: - audio_features = self._ensure_prompt_feature_batch(audio_features, "audio prompt embeds").to(device=self.device) - mask = self._ensure_prompt_mask_batch(mask).to(device=self.device) - additive_mask = convert_to_additive_mask(mask, video_features.dtype) - with torch.no_grad(): - video_embeds, audio_embeds, attention_mask = self.embeddings_processor.create_embeddings( - video_features, - audio_features, - additive_mask, - ) - return video_embeds.to(dtype=self.running_dtype), audio_embeds.to(dtype=self.running_dtype), attention_mask - - @staticmethod - def _ensure_prompt_feature_batch(value, name): - if not torch.is_tensor(value): - raise TypeError(f"LTX2 {name} must be a tensor, got {type(value)!r}.") - if value.ndim == 2: - return value.unsqueeze(0) - if value.ndim == 3: - return value - raise ValueError(f"LTX2 {name} must have shape [seq, dim] or [batch, seq, dim], got {tuple(value.shape)}.") - - @staticmethod - def _ensure_prompt_mask_batch(value): - if not torch.is_tensor(value): - raise TypeError(f"LTX2 prompt_attention_mask must be a tensor, got {type(value)!r}.") - if value.ndim == 1: - return value.unsqueeze(0) - if value.ndim == 2: - return value - raise ValueError(f"LTX2 prompt_attention_mask must have shape [seq] or [batch, seq], got {tuple(value.shape)}.") - - def encode_prompt_condition(self, prompt): - if self.text_encoder is None: - raise RuntimeError("LTX2 text encoder is not loaded. Set model.load_text_encoder=True for prompt-based training.") - if self.embeddings_processor is None: - raise RuntimeError("LTX2 embeddings processor is not loaded. Set model.load_embeddings_processor=True.") - if getattr(self.embeddings_processor, "feature_extractor", None) is None: - raise RuntimeError("LTX2 embeddings processor feature_extractor is required for prompt-based training.") - - prompts = [prompt] if isinstance(prompt, str) else list(prompt) - with torch.no_grad(): - encoded = self.text_encoder.encode(prompts, padding_side="left") - - hidden_states = [] - for layer_idx in range(len(encoded[0][0])): - layer = torch.cat([item[0][layer_idx] for item in encoded], dim=0) - hidden_states.append(layer.to(device=self.device, dtype=self.embeddings_dtype)) - attention_mask = torch.cat([item[1] for item in encoded], dim=0).to(device=self.device) - - with torch.no_grad(): - processed = self.embeddings_processor.process_hidden_states( - tuple(hidden_states), - attention_mask, - padding_side="left", - ) - audio_encoding = processed.audio_encoding - if audio_encoding is None: - audio_encoding = processed.video_encoding - return ( - processed.video_encoding.to(dtype=self.running_dtype), - audio_encoding.to(dtype=self.running_dtype), - processed.attention_mask, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/__init__.py new file mode 100644 index 000000000..97fdfe9b0 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/__init__.py @@ -0,0 +1,5 @@ +"""MiniMax-H3 model wrappers and training capabilities.""" + +from .minimax_h3_t2av import MiniMaxH3T2AVModel + +__all__ = ["MiniMaxH3T2AVModel"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/__init__.py new file mode 100644 index 000000000..7b1d10824 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/__init__.py @@ -0,0 +1,12 @@ +"""Capability adapters for MiniMax-H3.""" + +from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape +from .minimax_h3_distribution_matching_capability import MiniMaxH3DistributionMatchingCapability +from .minimax_h3_flow_matching_capability import MiniMaxH3FlowMatchingCapability + +__all__ = [ + "MiniMaxH3DistributionMatchingCapability", + "MiniMaxH3FlowMatchingCapability", + "MiniMaxH3JointLatents", + "MiniMaxH3LatentShape", +] diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/common.py new file mode 100644 index 000000000..574f1e1e5 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/common.py @@ -0,0 +1,27 @@ +"""Joint audio/video latent value objects used by MiniMax-H3.""" + +from dataclasses import dataclass + +from torch import Tensor + + +@dataclass(frozen=True) +class MiniMaxH3LatentShape: + """Packed latent geometry for one H3 audio/video sample.""" + + num_frames: int + latent_frames: int + latent_height: int + latent_width: int + audio_latents: int + video_tokens: tuple[int, int, int] + audio_tokens: tuple[int, int, int] + + +@dataclass(frozen=True) +class MiniMaxH3JointLatents: + """Video and audio token tensors that form one H3 diffusion state.""" + + video: Tensor + audio: Tensor + shape: MiniMaxH3LatentShape diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py new file mode 100644 index 000000000..7d72a4803 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py @@ -0,0 +1,377 @@ +"""Distribution-matching capability for MiniMax-H3 T2AV.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_capabilities import DistributionMatchingProfile +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _cached_condition, + _require_single_prompt, + _require_singleton_tensor, +) +from lightx2v_train.model_zoo.native.minimax_h3 import ( + audio_latent_num_frames, + build_packed_sequence, + build_row_timesteps, + video_latent_num_frames, +) + +from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape + + +@dataclass(frozen=True) +class MiniMaxH3DistributionMatchingOptions: + video_loss_weight: float = 1.0 + audio_loss_weight: float = 1.0 + video_flow_shift: float = 6.0 + audio_flow_shift: float = 3.0 + + @classmethod + def from_mapping(cls, config: Mapping | None) -> "MiniMaxH3DistributionMatchingOptions": + if config is None: + config = {} + if not isinstance(config, Mapping): + raise ValueError("model.capabilities.distribution_matching must be a mapping.") + options = cls( + video_loss_weight=float(config.get("video_loss_weight", 1.0)), + audio_loss_weight=float(config.get("audio_loss_weight", 1.0)), + video_flow_shift=float(config.get("video_flow_shift", 6.0)), + audio_flow_shift=float(config.get("audio_flow_shift", 3.0)), + ) + if options.video_flow_shift <= 0 or options.audio_flow_shift <= 0: + raise ValueError("MiniMax-H3 video_flow_shift and audio_flow_shift must be positive.") + if options.video_loss_weight < 0 or options.audio_loss_weight < 0: + raise ValueError("MiniMax-H3 video_loss_weight and audio_loss_weight cannot be negative.") + if options.video_loss_weight == 0 and options.audio_loss_weight == 0: + raise ValueError("At least one MiniMax-H3 modality loss weight must be non-zero.") + return options + + +def _shift_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor: + """Apply H3's rational flow shift to an unshifted noise level.""" + + return shift * sigma / (1.0 + (shift - 1.0) * sigma) + + +def _expand_sigma(sigma: torch.Tensor, ndim: int) -> torch.Tensor: + if sigma.ndim == 0: + sigma = sigma.reshape(1) + return sigma.reshape(sigma.shape[0], *([1] * (ndim - 1))) + + +class MiniMaxH3DistributionMatchingCapability(GenericDistributionMatchingCapability): + """H3-specific operations consumed by the framework's generic DMD loop. + + H3 jointly denoises packed video and stereo-audio tokens. It also uses a + clean-ward velocity (``x0 - noise``), while most LightX2V models use the + opposite flow direction. Keeping those details here lets ``DmdTrainer`` + manage roles, rollout, optimization, checkpointing, and FSDP unchanged. + """ + + _DEFAULT_LORA_TARGETS = ( + "to_q", + "to_k", + "to_v", + "to_out.0", + "ff.net.0.proj", + "ff.net.2", + ) + _PROFILE = DistributionMatchingProfile( + supported_training_methods=frozenset({"dmd"}), + supports_guidance=False, + supports_ida=False, + supports_diversity=False, + supports_real_data_fake=False, + supports_warped_denoising_schedule=False, + default_latent_dtype=torch.float32, + ) + + def __init__(self, model, options: Mapping | None = None) -> None: + super().__init__(model) + options = MiniMaxH3DistributionMatchingOptions.from_mapping(options) + self.video_weight = options.video_loss_weight + self.audio_weight = options.audio_loss_weight + self.video_shift = options.video_flow_shift + self.audio_shift = options.audio_flow_shift + self._layout_cache = {} + + @property + def profile(self) -> DistributionMatchingProfile: + return self._PROFILE + + @property + def default_negative_prompt(self): + return "" + + @property + def default_lora_target_modules(self): + return self._DEFAULT_LORA_TARGETS + + def latent_shape( + self, + batch, + shape_config, + image_sizes, + broadcast, + ): + if image_sizes: + raise ValueError("MiniMax-H3 uses training.dmd.height/width and does not support training.dmd.image_sizes.") + prompt = batch["conditioning"].get("prompt", "") + _require_single_prompt(prompt) + try: + height = int(broadcast(int(shape_config["height"]))) + width = int(broadcast(int(shape_config["width"]))) + num_frames = int(broadcast(int(shape_config["num_frames"]))) + except KeyError as exc: + raise KeyError(f"MiniMax-H3 DMD requires training.dmd.{exc.args[0]}.") from exc + + if height <= 0 or width <= 0 or num_frames <= 0: + raise ValueError(f"MiniMax-H3 dimensions must be positive, got frames={num_frames}, size={height}x{width}.") + size_multiple = self.model.vae_spatial_scale_factor * self.model.patch_size[1] + width_multiple = self.model.vae_spatial_scale_factor * self.model.patch_size[2] + if height % size_multiple or width % width_multiple: + raise ValueError(f"MiniMax-H3 height/width must be divisible by the VAE and patch scales ({size_multiple}, {width_multiple}), got {height}x{width}.") + + latent_frames = video_latent_num_frames(num_frames) + latent_height = height // self.model.vae_spatial_scale_factor + latent_width = width // self.model.vae_spatial_scale_factor + patch_t, patch_h, patch_w = self.model.patch_size + if patch_t != 1: + raise ValueError(f"MiniMax-H3 T2AV requires temporal patch size 1, got {self.model.patch_size}.") + video_rows = latent_frames * (latent_height // patch_h) * (latent_width // patch_w) + video_dimension = self.model.video_latent_channels * patch_t * patch_h * patch_w + audio_latents = audio_latent_num_frames(num_frames) + return MiniMaxH3LatentShape( + num_frames=num_frames, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + audio_latents=audio_latents, + video_tokens=(1, video_rows, video_dimension), + audio_tokens=(1, audio_latents * 2, self.model.audio_latent_channels), + ) + + def encode_conditions( + self, + batch, + negative_prompt, + guidance_scale, + broadcast, + ): + del negative_prompt + if guidance_scale != 1.0: + raise ValueError("MiniMax-H3 is guidance-distilled; training.teacher.guidance_scale must be 1.0.") + cached_condition = _cached_condition(batch, self.model) + with torch.no_grad(): + condition = self.model.encode_condition(batch) if cached_condition is None else self.model.prepare_text_condition(cached_condition) + return broadcast(condition), None + + def predict_velocity(self, latents, sigma, condition): + self._validate_latents(latents) + video_sigma, audio_sigma = self._modality_sigmas(sigma) + layout = self._layout(condition, latents.shape) + timesteps, timestep_indices = build_row_timesteps( + layout, + video_sigma, + audio_sigma, + ) + with self.model.transformer_forward_context(): + prediction = self.model.denoiser_module()( + hidden_states=latents.video, + audio_hidden_states=latents.audio, + encoder_hidden_states=condition["prompt_embeds"], + timestep=timesteps.to(self.device), + timestep_indices=timestep_indices.to(self.device), + token_tags=layout.token_tags, + position_ids=layout.position_ids, + video_indices=layout.video_indices, + audio_indices=layout.audio_indices, + text_indices=layout.text_indices, + return_dict=False, + ) + if not isinstance(prediction, (tuple, list)) or len(prediction) < 2: + raise TypeError("MiniMax-H3 transformer must return (video_velocity, audio_velocity) when return_dict=False.") + return MiniMaxH3JointLatents( + video=prediction[0], + audio=prediction[1], + shape=latents.shape, + ) + + def predict_guided_velocity( + self, + latents, + sigma, + condition, + negative_condition, + guidance_scale, + cfg_norm, + ): + del cfg_norm + if negative_condition is not None or guidance_scale != 1.0: + raise ValueError("MiniMax-H3 has no unconditional branch and only supports guidance_scale=1.0.") + return self.predict_velocity(latents, sigma, condition) + + def initial_latents(self, latent_shape, dtype, broadcast): + if latent_shape.video_tokens[0] != 1 or latent_shape.audio_tokens[0] != 1: + raise ValueError(f"MiniMax-H3 DMD requires physical batch size 1, got {latent_shape}.") + video = broadcast(torch.randn(latent_shape.video_tokens, device=self.device, dtype=dtype)) + audio = broadcast(torch.randn(latent_shape.audio_tokens, device=self.device, dtype=dtype)) + return MiniMaxH3JointLatents(video, audio, latent_shape) + + @staticmethod + def latent_hw(latent_shape): + del latent_shape + return None + + @staticmethod + def random_noise_like(latents, dtype, broadcast): + return MiniMaxH3JointLatents( + broadcast(torch.randn_like(latents.video, dtype=dtype)), + broadcast(torch.randn_like(latents.audio, dtype=dtype)), + latents.shape, + ) + + def add_noise(self, scheduler, latents, noise, sigma): + del scheduler + video_sigma, audio_sigma = self._modality_sigmas(sigma) + return MiniMaxH3JointLatents( + self._mix_noise(latents.video, noise.video, video_sigma), + self._mix_noise(latents.audio, noise.audio, audio_sigma), + latents.shape, + ) + + @staticmethod + def training_target(latents, noise): + # H3's time is t=1-sigma, so its velocity points from noise to x0. + return MiniMaxH3JointLatents( + latents.video.float() - noise.video.float(), + latents.audio.float() - noise.audio.float(), + latents.shape, + ) + + def step(self, scheduler, velocity, step_index, sample): + sigma = scheduler.sigma_at(step_index, device=self.device, dtype=torch.float32) + sigma_next = scheduler.sigma_at(int(step_index) + 1, device=self.device, dtype=torch.float32) + video_sigma, audio_sigma = self._modality_sigmas(sigma) + video_sigma_next, audio_sigma_next = self._modality_sigmas(sigma_next) + return ( + MiniMaxH3JointLatents( + self._cleanward_step(sample.video, velocity.video, video_sigma, video_sigma_next), + self._cleanward_step(sample.audio, velocity.audio, audio_sigma, audio_sigma_next), + sample.shape, + ), + MiniMaxH3JointLatents( + self._cleanward_x0(sample.video, velocity.video, video_sigma), + self._cleanward_x0(sample.audio, velocity.audio, audio_sigma), + sample.shape, + ), + ) + + def x0_from_velocity(self, sample, velocity, sigma): + video_sigma, audio_sigma = self._modality_sigmas(sigma) + return MiniMaxH3JointLatents( + self._cleanward_x0(sample.video, velocity.video, video_sigma), + self._cleanward_x0(sample.audio, velocity.audio, audio_sigma), + sample.shape, + ) + + def regression_loss(self, prediction, target): + video_loss = F.mse_loss(prediction.video.float(), target.video.float()) + audio_loss = F.mse_loss(prediction.audio.float(), target.audio.float()) + return self.video_weight * video_loss + self.audio_weight * audio_loss + + def dmd_loss(self, latents, fake_x0, teacher_x0): + return self.video_weight * super().dmd_loss( + latents.video, + fake_x0.video, + teacher_x0.video, + ) + self.audio_weight * super().dmd_loss( + latents.audio, + fake_x0.audio, + teacher_x0.audio, + ) + + @staticmethod + def detach(value): + return MiniMaxH3JointLatents( + value.video.detach(), + value.audio.detach(), + value.shape, + ) + + @staticmethod + def to_dtype(value, dtype): + return MiniMaxH3JointLatents( + value.video.to(dtype=dtype), + value.audio.to(dtype=dtype), + value.shape, + ) + + def extract_real_latents(self, batch, dtype, broadcast): + del batch, dtype, broadcast + raise ValueError("MiniMax-H3 joint DMD does not support real-data fake loss.") + + def _modality_sigmas(self, sigma): + sigma = torch.as_tensor(sigma, device=self.device, dtype=torch.float32) + if sigma.ndim == 0: + sigma = sigma.reshape(1) + if sigma.numel() != 1: + raise ValueError(f"MiniMax-H3 currently requires one shared base sigma, got shape {tuple(sigma.shape)}.") + return _shift_sigma(sigma, self.video_shift), _shift_sigma(sigma, self.audio_shift) + + def _layout(self, condition, shape): + tags = condition["text_token_tags"] + if tags.ndim != 1: + raise ValueError(f"MiniMax-H3 text_token_tags must be one-dimensional, got {tuple(tags.shape)}.") + if not bool((tags == 1).all()): + raise ValueError("MiniMax-H3 T2AV DMD currently supports text-only cached conditions.") + key = ( + int(tags.numel()), + shape.latent_frames, + shape.latent_height, + shape.latent_width, + shape.audio_latents, + self.model.patch_size, + self.device, + ) + layout = self._layout_cache.get(key) + if layout is None: + layout = build_packed_sequence( + tags.detach().cpu(), + shape.latent_frames, + shape.latent_height, + shape.latent_width, + shape.audio_latents, + self.model.patch_size, + ).to(self.device) + self._layout_cache[key] = layout + return layout + + @staticmethod + def _mix_noise(latent, noise, sigma): + expanded = _expand_sigma(sigma, latent.ndim) + return ((1.0 - expanded) * latent.float() + expanded * noise.float()).to(latent.dtype) + + @staticmethod + def _cleanward_step(sample, velocity, sigma, sigma_next): + current = _expand_sigma(sigma, sample.ndim) + following = _expand_sigma(sigma_next, sample.ndim) + return (sample.float() + (current - following) * velocity.float()).to(sample.dtype) + + @staticmethod + def _cleanward_x0(sample, velocity, sigma): + expanded = _expand_sigma(sigma, sample.ndim) + return (sample.float() + expanded * velocity.float()).to(sample.dtype) + + @staticmethod + def _validate_latents(latents): + if not isinstance(latents, MiniMaxH3JointLatents): + raise TypeError(f"Expected MiniMaxH3JointLatents, got {type(latents)!r}.") + _require_singleton_tensor(latents.video, "MiniMax-H3 video latent") + _require_singleton_tensor(latents.audio, "MiniMax-H3 audio latent") diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_flow_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_flow_matching_capability.py new file mode 100644 index 000000000..f19e2d104 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_flow_matching_capability.py @@ -0,0 +1,343 @@ +"""Flow-matching SFT for MiniMax-H3 joint video/audio latents.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_capabilities import ( + BoundCapability, + FlowMatchingSFTCapability, + LossResult, + SFTStepContext, +) +from lightx2v_train.model_zoo.capability_adapters.common import ( + _training_cache_data, + _uses_prompt_dropout, +) +from lightx2v_train.model_zoo.native.minimax_h3 import ( + audio_latent_num_frames, + build_packed_sequence, + build_row_timesteps, +) + +from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape + + +@dataclass(frozen=True) +class MiniMaxH3FlowMatchingOptions: + video_loss_weight: float = 1.0 + audio_loss_weight: float = 1.0 + video_flow_shift: float = 6.0 + audio_flow_shift: float = 3.0 + + @classmethod + def from_mapping(cls, config: Mapping | None) -> "MiniMaxH3FlowMatchingOptions": + if config is None: + config = {} + if not isinstance(config, Mapping): + raise ValueError("model.capabilities.flow_matching must be a mapping.") + options = cls( + video_loss_weight=float(config.get("video_loss_weight", 1.0)), + audio_loss_weight=float(config.get("audio_loss_weight", 1.0)), + video_flow_shift=float(config.get("video_flow_shift", 6.0)), + audio_flow_shift=float(config.get("audio_flow_shift", 3.0)), + ) + if options.video_flow_shift <= 0 or options.audio_flow_shift <= 0: + raise ValueError("MiniMax-H3 video_flow_shift and audio_flow_shift must be positive.") + if options.video_loss_weight < 0 or options.audio_loss_weight < 0: + raise ValueError("MiniMax-H3 video_loss_weight and audio_loss_weight cannot be negative.") + if options.video_loss_weight == 0 and options.audio_loss_weight == 0: + raise ValueError("At least one MiniMax-H3 modality loss weight must be non-zero.") + return options + + +class MiniMaxH3FlowMatchingCapability(BoundCapability, FlowMatchingSFTCapability): + """Expose H3's two-modality flow through the framework SFT contract.""" + + def __init__(self, model, options: Mapping | None = None) -> None: + super().__init__(model) + options = MiniMaxH3FlowMatchingOptions.from_mapping(options) + self.video_weight = options.video_loss_weight + self.audio_weight = options.audio_loss_weight + self.video_shift = options.video_flow_shift + self.audio_shift = options.audio_flow_shift + self._layout_cache = {} + + def encode_training_cache(self, batch): + with torch.no_grad(): + latents = self._joint_latents(batch) + prompt = batch["conditioning"]["prompt"] + prompts = {"positive": prompt} + if _uses_prompt_dropout(self.model): + prompts["unconditional"] = self.model.unconditional_prompt + return _training_cache_data( + self.model, + batch, + inputs={ + "video_latents": { + "tokens": latents.video, + "latent_frames": latents.shape.latent_frames, + "latent_height": latents.shape.latent_height, + "latent_width": latents.shape.latent_width, + }, + "audio_latents": latents.audio, + }, + prompts=prompts, + ) + + def compute_loss( + self, + batch: Mapping, + context: SFTStepContext, + ) -> LossResult: + scheduler = context.noise_scheduler + if getattr(scheduler, "do_time_shift", False): + raise ValueError("MiniMax-H3 SFT applies separate modality flow shifts; set scheduler.time_shift_settings.do_time_shift=false.") + + with torch.no_grad(): + latents = self._joint_latents(batch) + latents = MiniMaxH3JointLatents( + context.broadcast(latents.video), + context.broadcast(latents.audio), + latents.shape, + ) + condition = context.broadcast(self._condition(batch)) + sigma = context.broadcast(scheduler.sample_timestep_or_sigma()) + sigma = torch.as_tensor(sigma, device=self.model.device, dtype=torch.float32).reshape(-1) + if sigma.numel() != 1: + raise ValueError(f"MiniMax-H3 SFT requires one shared base sigma, got shape {tuple(sigma.shape)}.") + + noise = MiniMaxH3JointLatents( + context.broadcast(torch.randn_like(latents.video, dtype=self._latent_dtype)), + context.broadcast(torch.randn_like(latents.audio, dtype=self._latent_dtype)), + latents.shape, + ) + video_sigma, audio_sigma = self._modality_sigmas(sigma) + noisy = MiniMaxH3JointLatents( + self._mix_noise(latents.video, noise.video, video_sigma), + self._mix_noise(latents.audio, noise.audio, audio_sigma), + latents.shape, + ) + target = MiniMaxH3JointLatents( + latents.video.float() - noise.video.float(), + latents.audio.float() - noise.audio.float(), + latents.shape, + ) + + prediction = self._predict(noisy, condition, video_sigma, audio_sigma) + video_loss = F.mse_loss(prediction.video.float(), target.video) + audio_loss = F.mse_loss(prediction.audio.float(), target.audio) + loss = self.video_weight * video_loss + self.audio_weight * audio_loss + return LossResult( + loss=loss, + metrics={ + "video_loss": video_loss.detach(), + "audio_loss": audio_loss.detach(), + }, + ) + + @property + def _latent_dtype(self): + return getattr(self.model, "latent_dtype", torch.float32) + + def _condition(self, batch): + conditioning = batch.get("conditioning", {}) + active = conditioning.get("active", "positive") + cached = conditioning.get(active) + if cached is None: + condition = self.model.encode_condition(batch) + else: + condition = self.model.prepare_text_condition(cached) + if not bool((condition["text_token_tags"] == 1).all()): + raise ValueError("MiniMax-H3 T2AV SFT currently supports text-only cached conditions.") + return condition + + def _joint_latents(self, batch): + inputs = batch.get("inputs", {}) + video_value = inputs.get("video_latents") + if video_value is None: + video_value = inputs.get("latents") + audio_value = inputs.get("audio_latents") + if video_value is None and audio_value is None: + encoded = self.model.encode_to_cache_latents(batch) + video_value = encoded["video_latents"] + audio_value = encoded["audio_latents"] + elif video_value is None or audio_value is None: + raise KeyError("MiniMax-H3 requires both inputs.video_latents and inputs.audio_latents.") + + video, latent_frames, latent_height, latent_width = self._video_tokens(video_value) + audio = self._audio_tokens(audio_value) + audio_latents = audio.shape[1] // 2 + num_frames = self._source_num_frames(latent_frames) + expected_audio_latents = audio_latent_num_frames(num_frames) + if audio_latents != expected_audio_latents: + raise ValueError(f"MiniMax-H3 video geometry implies {expected_audio_latents} audio latents for {num_frames} source frames, but the cache contains {audio_latents}.") + shape = MiniMaxH3LatentShape( + num_frames=num_frames, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + audio_latents=audio_latents, + video_tokens=tuple(video.shape), + audio_tokens=tuple(audio.shape), + ) + return MiniMaxH3JointLatents(video.contiguous(), audio.contiguous(), shape) + + def _video_tokens(self, value): + metadata = value if isinstance(value, Mapping) else {} + tensor = self._cached_tensor(value, "video_latents") + if tensor.ndim == 2: + tensor = tensor.unsqueeze(0) + if tensor.ndim == 3: + latent_frames = self._metadata_int(metadata, "latent_frames") + latent_height = self._metadata_int(metadata, "latent_height") + latent_width = self._metadata_int(metadata, "latent_width") + tokens = tensor.to(device=self.model.device, dtype=self._latent_dtype) + elif tensor.ndim in {4, 5}: + if tensor.ndim == 4: + tensor = tensor.unsqueeze(0) + if tensor.shape[0] != 1 or tensor.shape[1] != self.model.video_latent_channels: + raise ValueError(f"MiniMax-H3 raw video latents must have shape [1,{self.model.video_latent_channels},F,H,W], got {tuple(tensor.shape)}.") + tensor = tensor.to(device=self.model.device, dtype=self._latent_dtype) + latent_frames, latent_height, latent_width = map(int, tensor.shape[-3:]) + tokens = self._patchify_video(tensor) + else: + raise ValueError(f"MiniMax-H3 video latents must be patchified [B,N,D] or raw [B,C,F,H,W], got {tuple(tensor.shape)}.") + + expected_dimension = self.model.video_latent_channels * math.prod(self.model.patch_size) + patch_t, patch_h, patch_w = self.model.patch_size + if latent_frames % patch_t or latent_height % patch_h or latent_width % patch_w: + raise ValueError(f"MiniMax-H3 video latent geometry {(latent_frames, latent_height, latent_width)} is not divisible by patch size {self.model.patch_size}.") + expected_rows = (latent_frames // patch_t) * (latent_height // patch_h) * (latent_width // patch_w) + if tokens.shape != (1, expected_rows, expected_dimension): + raise ValueError(f"MiniMax-H3 video tokens must have shape {(1, expected_rows, expected_dimension)}, got {tuple(tokens.shape)}.") + return tokens, latent_frames, latent_height, latent_width + + def _audio_tokens(self, value): + tensor = self._cached_tensor(value, "audio_latents") + if tensor.ndim == 2: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 3 or tensor.shape[0] != 1 or tensor.shape[-1] != self.model.audio_latent_channels: + raise ValueError(f"MiniMax-H3 audio latents must have shape [1,N,{self.model.audio_latent_channels}], got {tuple(tensor.shape)}.") + if tensor.shape[1] % 2: + raise ValueError(f"MiniMax-H3 stereo audio requires an even token count, got {tensor.shape[1]}.") + return tensor.to(device=self.model.device, dtype=self._latent_dtype) + + def _predict(self, latents, condition, video_sigma, audio_sigma): + layout = self._layout(condition, latents.shape) + timesteps, timestep_indices = build_row_timesteps(layout, video_sigma, audio_sigma) + with self.model.transformer_forward_context(): + output = self.model.denoiser_module()( + hidden_states=latents.video, + audio_hidden_states=latents.audio, + encoder_hidden_states=condition["prompt_embeds"], + timestep=timesteps.to(self.model.device), + timestep_indices=timestep_indices.to(self.model.device), + token_tags=layout.token_tags, + position_ids=layout.position_ids, + video_indices=layout.video_indices, + audio_indices=layout.audio_indices, + text_indices=layout.text_indices, + return_dict=False, + ) + if not isinstance(output, (tuple, list)) or len(output) < 2: + raise TypeError("MiniMax-H3 transformer must return (video_velocity, audio_velocity) when return_dict=False.") + if output[0].shape != latents.video.shape or output[1].shape != latents.audio.shape: + raise ValueError( + "MiniMax-H3 transformer output shapes do not match the SFT latents: " + f"video={tuple(output[0].shape)} vs {tuple(latents.video.shape)}, " + f"audio={tuple(output[1].shape)} vs {tuple(latents.audio.shape)}." + ) + return MiniMaxH3JointLatents(output[0], output[1], latents.shape) + + def _layout(self, condition, shape): + tags = condition["text_token_tags"] + key = ( + int(tags.numel()), + shape.latent_frames, + shape.latent_height, + shape.latent_width, + shape.audio_latents, + self.model.patch_size, + self.model.device, + ) + layout = self._layout_cache.get(key) + if layout is None: + layout = build_packed_sequence( + tags.detach().cpu(), + shape.latent_frames, + shape.latent_height, + shape.latent_width, + shape.audio_latents, + self.model.patch_size, + ).to(self.model.device) + self._layout_cache[key] = layout + return layout + + def _modality_sigmas(self, sigma): + return self._shift_sigma(sigma, self.video_shift), self._shift_sigma(sigma, self.audio_shift) + + @staticmethod + def _shift_sigma(sigma, shift): + return shift * sigma / (1.0 + (shift - 1.0) * sigma) + + @staticmethod + def _mix_noise(latent, noise, sigma): + sigma = sigma.reshape(sigma.shape[0], *([1] * (latent.ndim - 1))) + return ((1.0 - sigma) * latent.float() + sigma * noise.float()).to(latent.dtype) + + def _patchify_video(self, video): + batch, channels, frames, height, width = video.shape + patch_t, patch_h, patch_w = self.model.patch_size + return ( + video.reshape( + batch, + channels, + frames // patch_t, + patch_t, + height // patch_h, + patch_h, + width // patch_w, + patch_w, + ) + .permute(0, 2, 4, 6, 1, 3, 5, 7) + .reshape(batch, -1, channels * patch_t * patch_h * patch_w) + ) + + @staticmethod + def _cached_tensor(value, name): + if isinstance(value, Mapping): + value = value.get("tokens", value.get("latents")) + if not torch.is_tensor(value): + raise TypeError(f"MiniMax-H3 {name} cache must contain a tensor, got {type(value)!r}.") + return value + + @staticmethod + def _metadata_int(metadata, key): + value = metadata.get(key) + if value is None: + raise KeyError(f"Patchified MiniMax-H3 video latents require {key} metadata.") + if torch.is_tensor(value): + if value.numel() != 1: + raise ValueError(f"MiniMax-H3 {key} must be scalar, got shape {tuple(value.shape)}.") + value = value.item() + elif isinstance(value, (list, tuple)): + if len(value) != 1: + raise ValueError(f"MiniMax-H3 {key} must contain one value, got {len(value)}.") + value = value[0] + value = int(value) + if value <= 0: + raise ValueError(f"MiniMax-H3 {key} must be positive, got {value}.") + return value + + @staticmethod + def _source_num_frames(latent_frames): + chunk_latents = latent_frames - 2 + if chunk_latents < 0 or chunk_latents % 5: + raise ValueError(f"MiniMax-H3 latent frame count must be 5*n+2, got {latent_frames}.") + return chunk_latents // 5 * 17 + 5 diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/condition_encoder.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/condition_encoder.py new file mode 100644 index 000000000..59c1500b2 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/condition_encoder.py @@ -0,0 +1,141 @@ +"""Diffusers/Transformers conditioner adapter for MiniMax-H3 training.""" + +import torch + + +class MiniMaxH3ConditionEncoder: + """Encode text with the official Qwen3-VL conditioning recipe.""" + + def __init__( + self, + model_path, + *, + device, + dtype, + local_files_only=True, + cpu_offload=False, + attention_backend="torch_sdpa", + text_encoder_subfolder="text_encoder", + tokenizer_subfolder="tokenizer", + processor_subfolder="processor", + text_encoder_layer=50, + text_tag=1, + ): + try: + from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor + except ImportError as error: + raise ImportError("MiniMax-H3 cache construction requires Transformers with Qwen3-VL support.") from error + + self.device = torch.device(device) + self.dtype = dtype + self.cpu_offload = bool(cpu_offload) + self._managed_cpu_offload = False + self.text_encoder_layer = int(text_encoder_layer) + self.text_tag = int(text_tag) + if self.text_encoder_layer < 0: + raise ValueError("MiniMax-H3 text_encoder_layer must be non-negative.") + + load_kwargs = { + "local_files_only": bool(local_files_only), + "attn_implementation": self._transformers_attention_backend(attention_backend), + "low_cpu_mem_usage": True, + } + try: + self.encoder = Qwen3VLForConditionalGeneration.from_pretrained( + model_path, + subfolder=text_encoder_subfolder, + dtype=dtype, + **load_kwargs, + ) + except TypeError: + self.encoder = Qwen3VLForConditionalGeneration.from_pretrained( + model_path, + subfolder=text_encoder_subfolder, + torch_dtype=dtype, + **load_kwargs, + ) + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + model_path, + subfolder=tokenizer_subfolder, + local_files_only=bool(local_files_only), + ) + self.processor = Qwen3VLProcessor.from_pretrained( + model_path, + subfolder=processor_subfolder, + local_files_only=bool(local_files_only), + ) + if not hasattr(self.processor, "create_mm_token_type_ids"): + raise RuntimeError("Installed Transformers does not expose Qwen3VLProcessor.create_mm_token_type_ids required by MiniMax-H3.") + + self.encoder.requires_grad_(False).eval() + if self.cpu_offload and self.device.type != "cpu": + try: + from accelerate import cpu_offload + except ImportError as error: + raise ImportError("MiniMax-H3 condition encoder CPU offload requires Accelerate.") from error + cpu_offload(self.encoder, execution_device=self.device, offload_buffers=True) + self._managed_cpu_offload = True + else: + self.encoder.to(self.device) + + @staticmethod + def _transformers_attention_backend(attention_backend): + if attention_backend in {None, "torch_sdpa", "native"}: + return "sdpa" + return attention_backend + + def _activate(self): + if self.cpu_offload and self.device.type != "cpu" and not self._managed_cpu_offload: + self.encoder.to(self.device) + + def _offload(self): + if self.cpu_offload and self.device.type != "cpu" and not self._managed_cpu_offload: + self.encoder.to("cpu") + + @torch.inference_mode() + def encode(self, prompt): + if not isinstance(prompt, str): + prompts = list(prompt) + if len(prompts) != 1: + raise ValueError(f"MiniMax-H3 requires one prompt per rank, got {len(prompts)}.") + prompt = prompts[0] + if not isinstance(prompt, str): + raise TypeError(f"MiniMax-H3 prompt must be a string, got {type(prompt).__name__}.") + if not prompt: + raise ValueError("MiniMax-H3 prompt must contain at least one character.") + + token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"] + if not token_ids: + raise ValueError("MiniMax-H3 prompt must produce at least one token.") + self._activate() + try: + input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device) + mm_token_type_ids = torch.tensor( + self.processor.create_mm_token_type_ids([token_ids]), + dtype=torch.long, + device=self.device, + ) + outputs = self.encoder.model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + mm_token_type_ids=mm_token_type_ids, + use_cache=False, + output_hidden_states=True, + ) + if len(outputs.hidden_states) <= self.text_encoder_layer: + raise ValueError(f"MiniMax-H3 requires Qwen3-VL hidden_states[{self.text_encoder_layer}], but the encoder returned only {len(outputs.hidden_states)} states.") + prompt_embeds = outputs.hidden_states[self.text_encoder_layer].to(device=self.device, dtype=self.dtype) + del outputs + if prompt_embeds.shape[1] != len(token_ids): + raise ValueError(f"MiniMax-H3 Qwen3-VL returned {prompt_embeds.shape[1]} embedding rows for {len(token_ids)} tokens.") + return { + "prompt_embeds": prompt_embeds, + "text_token_tags": torch.full( + (len(token_ids),), + self.text_tag, + dtype=torch.long, + device=self.device, + ), + } + finally: + self._offload() diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/data_process.py new file mode 100644 index 000000000..732ffd426 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/data_process.py @@ -0,0 +1,98 @@ +"""Raw T2AV sample preparation for MiniMax-H3 cache construction.""" + +from pathlib import Path + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_zoo.native.minimax_h3 import audio_latent_num_frames, video_latent_num_frames +from lightx2v_train.utils.registry import SAMPLE_PROCESSOR_REGISTER + + +class MiniMaxH3T2AVProcessor: + """Prepare synchronized RGB video and stereo audio for the H3 VAEs.""" + + unconditional_prompt = " " + requires_audio = True + + def __init__(self, config): + model_config = config["model"] + self.video_fps = int(model_config.get("video_fps", 24)) + self.audio_sample_rate = int(model_config.get("audio_sampling_rate", 32000)) + self.audio_latents_per_second = int(model_config.get("audio_latents_per_second", 40)) + if self.video_fps <= 0 or self.audio_sample_rate <= 0 or self.audio_latents_per_second <= 0: + raise ValueError("MiniMax-H3 video/audio rates must be positive.") + if self.video_fps != 24 or self.audio_latents_per_second != 40: + raise ValueError("MiniMax-H3 T2AV requires video_fps=24 and audio_latents_per_second=40.") + if self.audio_sample_rate % self.audio_latents_per_second: + raise ValueError("MiniMax-H3 audio_sampling_rate must be divisible by audio_latents_per_second.") + self.audio_hop_length = self.audio_sample_rate // self.audio_latents_per_second + + def __call__(self, sample): + video = sample.get("inputs", {}).get("video") + if not torch.is_tensor(video) or video.ndim != 4 or video.shape[0] != 3: + shape = tuple(video.shape) if torch.is_tensor(video) else type(video).__name__ + raise ValueError(f"MiniMax-H3 expects decoded video [3,F,H,W], got {shape}.") + num_frames = int(video.shape[1]) + video_latent_num_frames(num_frames) + + audio_path = sample.get("meta", {}).get("audio_path") + if not audio_path: + raise ValueError("MiniMax-H3 T2AV cache construction requires an audio path for every video.") + waveform = self._load_audio(audio_path) + audio_frames = audio_latent_num_frames(num_frames) + target_samples = audio_frames * self.audio_hop_length + video_start_time = float(sample["meta"].get("video_start_time", 0.0)) + if video_start_time < 0: + raise ValueError(f"MiniMax-H3 video_start_time must be non-negative, got {video_start_time}.") + audio_start_sample = int(round(video_start_time * self.audio_sample_rate)) + waveform = self._slice_audio(waveform, audio_start_sample, target_samples) + + # VideoDataset normalizes pixels to [-1, 1], while H3's ImageNet VAE + # normalization consumes [0, 1]. + sample["inputs"]["video"] = ((video + 1.0) * 0.5).clamp_(0.0, 1.0) + sample["inputs"]["audio"] = waveform + sample["meta"].update( + { + "num_frames": num_frames, + "target_height": int(video.shape[-2]), + "target_width": int(video.shape[-1]), + "audio_sample_rate": self.audio_sample_rate, + "audio_start_sample": audio_start_sample, + } + ) + return sample + + def _load_audio(self, audio_path): + try: + import torchaudio + except ImportError as error: + raise ImportError("MiniMax-H3 cache construction requires torchaudio.") from error + + waveform, sample_rate = torchaudio.load(Path(audio_path)) + waveform = waveform.float() + if sample_rate != self.audio_sample_rate: + waveform = torchaudio.functional.resample( + waveform, + orig_freq=sample_rate, + new_freq=self.audio_sample_rate, + ) + if waveform.shape[0] == 1: + waveform = waveform.expand(2, -1) + elif waveform.shape[0] > 2: + waveform = waveform[:2] + if waveform.shape[0] != 2: + raise ValueError(f"MiniMax-H3 audio must be mono or stereo, got {waveform.shape[0]} channels.") + return waveform.contiguous() + + @staticmethod + def _slice_audio(waveform, start_sample, target_samples): + waveform = waveform[..., start_sample : start_sample + target_samples] + if waveform.shape[-1] < target_samples: + waveform = F.pad(waveform, (0, target_samples - waveform.shape[-1])) + return waveform.contiguous() + + +@SAMPLE_PROCESSOR_REGISTER("minimax_h3_t2av") +def build_minimax_h3_t2av_processor(config): + return MiniMaxH3T2AVProcessor(config) diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py new file mode 100644 index 000000000..5e7e6541d --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py @@ -0,0 +1,354 @@ +"""LightX2V-Train wrapper for the trainable MiniMax-H3 T2AV DiT.""" + +from collections.abc import Mapping +from contextlib import contextmanager, nullcontext +from pathlib import Path + +import torch +from peft import LoraConfig, inject_adapter_in_model + +from lightx2v_train.model_capabilities import ( + DistributionMatchingCapability, + FlowMatchingSFTCapability, +) +from lightx2v_train.model_zoo.minimax_h3.capability_adapters import ( + MiniMaxH3DistributionMatchingCapability, + MiniMaxH3FlowMatchingCapability, +) +from lightx2v_train.model_zoo.native.minimax_h3 import load_minimax_h3_transformer +from lightx2v_train.utils.registry import MODEL_REGISTER +from lightx2v_train.utils.utils import get_running_dtype + +from ..base import BaseModel +from .condition_encoder import MiniMaxH3ConditionEncoder + + +@MODEL_REGISTER("minimax_h3_t2av") +class MiniMaxH3T2AVModel(BaseModel): + """A standard trainable wrapper around Diffusers' MiniMax-H3 module.""" + + pipeline_cls = None + + def register_capabilities(self): + super().register_capabilities() + capability_config = self.config["model"].get("capabilities", {}) + if not isinstance(capability_config, Mapping): + raise ValueError("model.capabilities must be a mapping.") + if "distillation" in capability_config: + raise ValueError("model.capabilities.distillation was renamed to model.capabilities.distribution_matching.") + unsupported_capabilities = set(capability_config) - { + "flow_matching", + "distribution_matching", + } + if unsupported_capabilities: + names = ", ".join(sorted(unsupported_capabilities)) + raise ValueError(f"MiniMax-H3 currently supports only SFT (flow_matching) and DMD (distribution_matching); unsupported model capabilities: {names}.") + self.capabilities.register( + FlowMatchingSFTCapability, + MiniMaxH3FlowMatchingCapability( + self, + capability_config.get("flow_matching"), + ), + ) + self.capabilities.register( + DistributionMatchingCapability, + MiniMaxH3DistributionMatchingCapability( + self, + capability_config.get("distribution_matching"), + ), + ) + + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + config = self.config["model"] + self.pretrained_model_path = config["pretrained_model_name_or_path"] + self.transformer_param_dtype = get_running_dtype(config.get("transformer_param_dtype", "bf16")) + self.latent_dtype = get_running_dtype(config.get("latent_dtype", "fp32")) + self.patch_size = tuple(int(value) for value in config.get("patch_size", (1, 2, 2))) + if len(self.patch_size) != 3: + raise ValueError(f"model.patch_size must contain three integers, got {self.patch_size}.") + self.video_latent_channels = int(config.get("video_latent_channels", 24)) + self.audio_latent_channels = int(config.get("audio_latent_channels", 32)) + self.vae_spatial_scale_factor = int(config.get("vae_spatial_scale_factor", 16)) + self.audio_sampling_rate = int(config.get("audio_sampling_rate", 32000)) + if self.video_latent_channels <= 0 or self.audio_latent_channels <= 0 or self.vae_spatial_scale_factor <= 0 or self.audio_sampling_rate <= 0: + raise ValueError("MiniMax-H3 latent channels, VAE spatial scale, and audio rate must be positive.") + self.use_autocast = bool(config.get("use_autocast", False)) + self.cache_encoder_cpu_offload = bool(config.get("cache_encoder_cpu_offload", False)) + self.transformer = None + self.video_vae = None + self.audio_vae = None + self.condition_encoder = None + self._managed_cache_encoders = set() + if load_vae: + self._load_vaes(config) + if load_condition_encoder: + self.condition_encoder = MiniMaxH3ConditionEncoder( + self.pretrained_model_path, + device=self.device, + dtype=self.running_dtype, + local_files_only=bool(config.get("local_files_only", True)), + cpu_offload=bool( + config.get( + "cache_condition_encoder_cpu_offload", + config.get("cache_encoder_cpu_offload", False), + ) + ), + attention_backend=config.get( + "condition_attention_backend", + "torch_sdpa", + ), + text_encoder_subfolder=config.get("text_encoder_subfolder", "text_encoder"), + tokenizer_subfolder=config.get("tokenizer_subfolder", "tokenizer"), + processor_subfolder=config.get("processor_subfolder", "processor"), + text_encoder_layer=int(config.get("text_encoder_layer", 50)), + text_tag=int(config.get("text_tag", 1)), + ) + if load_transformer: + self.transformer = load_minimax_h3_transformer( + self.pretrained_model_path, + torch_dtype=self.transformer_param_dtype, + local_files_only=bool(config.get("local_files_only", True)), + attention_backend=config.get("attention_backend"), + ) + self.transformer.to(self.device) + + def _load_vaes(self, config): + try: + from diffusers import AutoencoderKLMiniMaxH3, AutoencoderKLMiniMaxH3Audio + except ImportError as error: + raise ImportError("MiniMax-H3 cache construction requires Diffusers with MiniMax-H3 video and audio VAEs.") from error + + local_files_only = bool(config.get("local_files_only", True)) + video_vae_subfolder = self._component_subfolder( + self.pretrained_model_path, + config.get("video_vae_subfolder"), + candidates=("vae", "video_vae"), + ) + self.video_vae = AutoencoderKLMiniMaxH3.from_pretrained( + self.pretrained_model_path, + subfolder=video_vae_subfolder, + torch_dtype=torch.float32, + local_files_only=local_files_only, + low_cpu_mem_usage=True, + ) + self.audio_vae = AutoencoderKLMiniMaxH3Audio.from_pretrained( + self.pretrained_model_path, + subfolder=config.get("audio_vae_subfolder", "audio_vae"), + torch_dtype=torch.float32, + local_files_only=local_files_only, + low_cpu_mem_usage=True, + ) + for vae in (self.video_vae, self.audio_vae): + vae.requires_grad_(False).eval() + + vae_attention_backend = config.get("vae_attention_backend") + if vae_attention_backend not in {None, "torch_sdpa", "sdpa", "native"}: + self.video_vae.set_attention_backend(vae_attention_backend) + if self.cache_encoder_cpu_offload and self.device.type != "cpu": + try: + from accelerate import cpu_offload + except ImportError as error: + raise ImportError("MiniMax-H3 cache encoder CPU offload requires Accelerate.") from error + self.video_vae.enable_group_offload( + onload_device=self.device, + offload_device=torch.device("cpu"), + offload_type="leaf_level", + ) + cpu_offload(self.audio_vae, execution_device=self.device, offload_buffers=True) + self._managed_cache_encoders.update({"video", "audio"}) + else: + self.video_vae.to(self.device) + self.audio_vae.to(self.device) + + video_channels = int(self.video_vae.config.latent_channels) + audio_channels = int(self.audio_vae.config.latent_channels) + sampling_rate = int(self.audio_vae.config.sampling_rate) + if video_channels != self.video_latent_channels: + raise ValueError(f"MiniMax-H3 video_latent_channels does not match the video VAE: config={self.video_latent_channels}, checkpoint={video_channels}.") + if audio_channels != self.audio_latent_channels: + raise ValueError(f"MiniMax-H3 audio_latent_channels does not match the audio VAE: config={self.audio_latent_channels}, checkpoint={audio_channels}.") + if sampling_rate != self.audio_sampling_rate: + raise ValueError(f"MiniMax-H3 audio_sampling_rate does not match the audio VAE: config={self.audio_sampling_rate}, checkpoint={sampling_rate}.") + + @staticmethod + def _component_subfolder(model_path, configured, candidates): + if configured: + return str(configured) + root = Path(str(model_path)).expanduser() + if root.is_dir(): + for candidate in candidates: + if (root / candidate / "config.json").is_file(): + return candidate + return candidates[0] + + @contextmanager + def _active_cache_encoder(self, encoder): + name = "video" if encoder is self.video_vae else "audio" if encoder is self.audio_vae else None + managed_offload = name in self._managed_cache_encoders + should_offload = self.cache_encoder_cpu_offload and self.device.type != "cpu" and not managed_offload + if should_offload: + encoder.to(self.device) + try: + yield encoder + finally: + if should_offload: + encoder.to("cpu") + + @staticmethod + def _latent_statistics(vae, latents, shape): + mean = latents.new_tensor(vae.config.latents_mean).view(shape) + std = latents.new_tensor(vae.config.latents_std).view(shape) + if bool((std == 0).any()): + raise ValueError("MiniMax-H3 VAE latent standard deviations must be non-zero.") + return mean, std + + @torch.inference_mode() + def _encode_video_latents(self, video): + with self._active_cache_encoder(self.video_vae) as video_vae: + pixels = video.to(device=self.device, dtype=torch.float32) + pixel_mean = pixels.new_tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1, 1) + pixel_std = pixels.new_tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1, 1) + posterior = video_vae.encode((pixels - pixel_mean) / pixel_std).latent_dist + latents = posterior.mode().float() + mean, std = self._latent_statistics(self.video_vae, latents, shape=(1, -1, 1, 1, 1)) + return (latents - mean) / std + + @torch.inference_mode() + def _encode_audio_latents(self, audio): + with self._active_cache_encoder(self.audio_vae) as audio_vae: + stereo_waveform = audio[0].unsqueeze(1).to(device=self.device, dtype=torch.float32) + posterior = audio_vae.encode(stereo_waveform).latent_dist + latents = posterior.mode().float().transpose(1, 2) + mean, std = self._latent_statistics(self.audio_vae, latents, shape=(1, 1, -1)) + latents = (latents - mean) / std + return latents.reshape(1, -1, self.audio_latent_channels) + + def reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self.video_vae = source.video_vae + self.audio_vae = source.audio_vae + self.condition_encoder = source.condition_encoder + self._managed_cache_encoders = set(source._managed_cache_encoders) + + def denoiser_module(self): + return self.transformer + + def transformer_forward_context(self): + if self.use_autocast and self.device.type == "cuda" and self.running_dtype in {torch.float16, torch.bfloat16}: + return torch.autocast("cuda", dtype=self.running_dtype) + return nullcontext() + + def add_lora(self, rank, alpha, target_modules): + if not target_modules: + target_modules = MiniMaxH3DistributionMatchingCapability._DEFAULT_LORA_TARGETS + lora_config = LoraConfig( + r=rank, + lora_alpha=alpha, + init_lora_weights="gaussian", + target_modules=target_modules, + ) + try: + self.transformer = inject_adapter_in_model( + lora_config, + self.transformer, + adapter_name="default", + ) + except TypeError: + self.transformer = inject_adapter_in_model(lora_config, self.transformer) + + def prepare_text_condition(self, condition): + if not isinstance(condition, dict): + raise TypeError(f"MiniMax-H3 cached condition must be a dict, got {type(condition)!r}.") + missing = {"prompt_embeds", "text_token_tags"} - condition.keys() + if missing: + names = ", ".join(sorted(missing)) + raise KeyError(f"MiniMax-H3 cached condition is missing: {names}.") + + prompt_embeds = condition["prompt_embeds"] + text_token_tags = condition["text_token_tags"] + if not torch.is_tensor(prompt_embeds) or not torch.is_tensor(text_token_tags): + raise TypeError("MiniMax-H3 prompt_embeds and text_token_tags must be tensors.") + if prompt_embeds.ndim == 2: + prompt_embeds = prompt_embeds.unsqueeze(0) + if prompt_embeds.ndim != 3 or prompt_embeds.shape[0] != 1: + raise ValueError(f"MiniMax-H3 prompt_embeds must have shape [1, tokens, dim], got {tuple(prompt_embeds.shape)}.") + if text_token_tags.ndim == 2: + if text_token_tags.shape[0] != 1: + raise ValueError("MiniMax-H3 currently requires data.train.batch_size=1.") + text_token_tags = text_token_tags[0] + if text_token_tags.ndim != 1 or text_token_tags.shape[0] != prompt_embeds.shape[1]: + raise ValueError(f"MiniMax-H3 text_token_tags must contain one tag per prompt embedding row; got {tuple(text_token_tags.shape)} for {prompt_embeds.shape[1]} rows.") + return { + "prompt_embeds": prompt_embeds.to(self.device, dtype=self.running_dtype), + "text_token_tags": text_token_tags.to(self.device, dtype=torch.long), + } + + def encode_prompt_condition(self, prompt): + if self.condition_encoder is None: + raise RuntimeError("MiniMax-H3 condition encoder is not loaded. Use a training cache or load condition components.") + return self.prepare_text_condition(self.condition_encoder.encode(prompt)) + + def encode_condition(self, sample): + conditioning = sample.get("conditioning", {}) + cached = conditioning.get(conditioning.get("active", "positive")) + if cached is not None: + return self.prepare_text_condition(cached) + return self.encode_prompt_condition(sample["conditioning"]["prompt"]) + + def encode_to_cache_latents(self, sample): + inputs = sample.get("inputs", {}) + video_latents = inputs.get("video_latents") + audio_latents = inputs.get("audio_latents") + if video_latents is not None and audio_latents is not None: + return { + "video_latents": video_latents, + "audio_latents": audio_latents, + } + if self.video_vae is None or self.audio_vae is None: + raise RuntimeError("MiniMax-H3 VAEs are not loaded. Use cached latents or load VAE components.") + + video = inputs.get("video") + audio = inputs.get("audio") + if not torch.is_tensor(video) or not torch.is_tensor(audio): + raise KeyError("MiniMax-H3 cache encoding requires inputs.video and inputs.audio.") + if video.ndim == 4: + video = video.unsqueeze(0) + if video.ndim != 5 or video.shape[:2] != (1, 3): + raise ValueError(f"MiniMax-H3 video must be [1,3,F,H,W], got {tuple(video.shape)}.") + if audio.ndim == 2: + audio = audio.unsqueeze(0) + if audio.ndim != 3 or audio.shape[:2] != (1, 2): + raise ValueError(f"MiniMax-H3 audio must be [1,2,samples], got {tuple(audio.shape)}.") + + video_latents = self._encode_video_latents(video) + audio_latents = self._encode_audio_latents(audio) + return { + "video_latents": video_latents, + "audio_latents": audio_latents, + } + + def enable_gradient_checkpointing(self): + if hasattr(self.transformer, "enable_gradient_checkpointing"): + self.transformer.enable_gradient_checkpointing() + else: + self.transformer.gradient_checkpointing = True + + def fsdp2_shard_plan(self, fsdp_config): + reshard = fsdp_config.get("reshard_after_forward", {}) + blocks = list(self.transformer.token_refiner.refiner_blocks) + list(self.transformer.transformer_blocks) + return [ + { + "modules": blocks, + "reshard_after_forward": reshard.get("block_reshard", True), + }, + { + "module": self.transformer, + "reshard_after_forward": reshard.get("root_reshard", False), + }, + ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3_t2av.py deleted file mode 100644 index f4eb6ca3c..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3_t2av.py +++ /dev/null @@ -1,99 +0,0 @@ -"""LightX2V-Train wrapper for the trainable MiniMax-H3 T2AV DiT.""" - -from contextlib import nullcontext - -import torch -from peft import LoraConfig, inject_adapter_in_model - -from lightx2v_train.utils.registry import MODEL_REGISTER -from lightx2v_train.utils.utils import get_running_dtype - -from .base import BaseModel -from .native.minimax_h3 import load_minimax_h3_transformer - - -@MODEL_REGISTER("minimax_h3_t2av") -class MiniMaxH3T2AVModel(BaseModel): - """A normal ``nn.Module`` H3 model suitable for LoRA/full DMD training.""" - - pipeline_cls = None - - def load_components(self, transformer_only=False, reference_model=None): - del transformer_only, reference_model - config = self.config["model"] - self.pretrained_model_path = config["pretrained_model_name_or_path"] - self.transformer_param_dtype = get_running_dtype(config.get("transformer_param_dtype", "bf16")) - self.patch_size = tuple(config.get("patch_size", (1, 2, 2))) - self.video_latent_channels = int(config.get("video_latent_channels", 24)) - self.audio_latent_channels = int(config.get("audio_latent_channels", 32)) - self.vae_spatial_scale_factor = int(config.get("vae_spatial_scale_factor", 16)) - self.use_autocast = bool(config.get("use_autocast", False)) - self.transformer = load_minimax_h3_transformer( - self.pretrained_model_path, - torch_dtype=self.transformer_param_dtype, - local_files_only=bool(config.get("local_files_only", True)), - attention_backend=config.get("attention_backend"), - ) - self.transformer.to(self.device) - - def denoiser_module(self): - return self.transformer - - def transformer_forward_context(self): - if self.use_autocast and self.device.type == "cuda" and self.running_dtype in {torch.float16, torch.bfloat16}: - return torch.autocast("cuda", dtype=self.running_dtype) - return nullcontext() - - def add_lora(self, rank, alpha, target_modules): - if not target_modules: - target_modules = ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2"] - lora_config = LoraConfig( - r=rank, - lora_alpha=alpha, - init_lora_weights="gaussian", - target_modules=target_modules, - ) - try: - self.transformer = inject_adapter_in_model(lora_config, self.transformer, adapter_name="default") - except TypeError: - self.transformer = inject_adapter_in_model(lora_config, self.transformer) - - def prepare_text_condition(self, condition): - if not isinstance(condition, dict): - raise TypeError(f"MiniMax-H3 cached condition must be a dict, got {type(condition)!r}.") - if "prompt_embeds" not in condition or "text_token_tags" not in condition: - raise KeyError("MiniMax-H3 condition requires prompt_embeds and text_token_tags.") - prompt_embeds = condition["prompt_embeds"] - text_token_tags = condition["text_token_tags"] - if prompt_embeds.ndim == 2: - prompt_embeds = prompt_embeds.unsqueeze(0) - if prompt_embeds.ndim != 3 or prompt_embeds.shape[0] != 1: - raise ValueError(f"MiniMax-H3 prompt_embeds must have shape [1, tokens, dim], got {tuple(prompt_embeds.shape)}.") - if text_token_tags.ndim == 2: - if text_token_tags.shape[0] != 1: - raise ValueError("MiniMax-H3 currently requires data.train.batch_size=1.") - text_token_tags = text_token_tags[0] - if text_token_tags.ndim != 1 or text_token_tags.shape[0] != prompt_embeds.shape[1]: - raise ValueError(f"MiniMax-H3 text_token_tags must contain one tag per prompt embedding row; got {tuple(text_token_tags.shape)} for {prompt_embeds.shape[1]} rows.") - return { - "prompt_embeds": prompt_embeds.to(self.device, dtype=self.running_dtype), - "text_token_tags": text_token_tags.to(self.device, dtype=torch.long), - } - - def encode_prompt_condition(self, prompt): - del prompt - raise RuntimeError("MiniMax-H3 DMD expects cached prompt latents. Run build_minimax_h3_prompt_latents.py first.") - - def enable_gradient_checkpointing(self): - if hasattr(self.transformer, "enable_gradient_checkpointing"): - self.transformer.enable_gradient_checkpointing() - else: - self.transformer.gradient_checkpointing = True - - def fsdp2_shard_plan(self, fsdp_config): - reshard = fsdp_config.get("reshard_after_forward", {}) - blocks = list(self.transformer.token_refiner.refiner_blocks) + list(self.transformer.transformer_blocks) - return [ - {"modules": blocks, "reshard_after_forward": reshard.get("block_reshard", True)}, - {"module": self.transformer, "reshard_after_forward": reshard.get("root_reshard", False)}, - ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/__init__.py deleted file mode 100644 index 2c5f945f7..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/__init__.py +++ /dev/null @@ -1,68 +0,0 @@ -import sys -from pathlib import Path - -_NATIVE_LTX2_ROOT = Path(__file__).resolve().parent - - -def ensure_ltx2_native_path(): - if not (_NATIVE_LTX2_ROOT / "ltx_core").is_dir(): - raise FileNotFoundError(f"Missing vendored LTX2 native package: {_NATIVE_LTX2_ROOT}") - - native_root = str(_NATIVE_LTX2_ROOT) - if native_root in sys.path: - sys.path.remove(native_root) - sys.path.insert(0, native_root) - - loaded = sys.modules.get("ltx_core") - if loaded is None: - return - - origin = Path(getattr(loaded, "__file__", "")).resolve() - try: - origin.relative_to(_NATIVE_LTX2_ROOT.resolve()) - except ValueError as exc: - raise RuntimeError(f"ltx_core is already loaded from a non-native path: {origin}") from exc - - -ensure_ltx2_native_path() - -from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier, get_pixel_coords # noqa: E402 -from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder # noqa: E402 -from ltx_core.model.transformer.modality import Modality # noqa: E402 -from ltx_core.model.transformer.model_configurator import ( # noqa: E402 - LTXV_MODEL_COMFY_RENAMING_MAP, - LTXModelConfigurator, -) -from ltx_core.text_encoders.gemma import ( # noqa: E402 - EMBEDDINGS_PROCESSOR_KEY_OPS, - GEMMA_LLM_KEY_OPS, - GEMMA_MODEL_OPS, - EmbeddingsProcessorConfigurator, - GemmaTextEncoderConfigurator, - convert_to_additive_mask, - module_ops_from_gemma_root, -) -from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape # noqa: E402 -from ltx_core.utils import find_matching_file # noqa: E402 - -__all__ = [ - "AudioLatentShape", - "AudioPatchifier", - "EMBEDDINGS_PROCESSOR_KEY_OPS", - "GEMMA_LLM_KEY_OPS", - "GEMMA_MODEL_OPS", - "EmbeddingsProcessorConfigurator", - "GemmaTextEncoderConfigurator", - "LTXModelConfigurator", - "LTXV_MODEL_COMFY_RENAMING_MAP", - "Modality", - "SingleGPUModelBuilder", - "SpatioTemporalScaleFactors", - "VideoLatentPatchifier", - "VideoLatentShape", - "convert_to_additive_mask", - "ensure_ltx2_native_path", - "find_matching_file", - "get_pixel_coords", - "module_ops_from_gemma_root", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/batch_split.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/batch_split.py deleted file mode 100644 index 9c6cbb0af..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/batch_split.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Batch-splitting adapter for the transformer. -Wraps an ``X0Model`` (or ``BlockStreamingWrapper``) and splits batched inputs -into smaller chunks before forwarding, then concatenates the results. This -controls peak activation memory at the cost of more forward passes. -The adapter is transparent — it has the same ``forward`` signature as -``X0Model`` and proxies attribute access to the wrapped model. -Example -------- ->>> from ltx_core.batch_split import BatchSplitAdapter ->>> adapter = BatchSplitAdapter(model, max_batch_size=1) ->>> # Receives B=4, runs 4xB=1 internally, returns B=4 ->>> denoised_video, denoised_audio = adapter(video=v_b4, audio=a_b4, perturbations=ptb) -""" - -from __future__ import annotations - -from typing import Any - -import torch -from ltx_core.guidance.perturbations import BatchedPerturbationConfig -from ltx_core.model.transformer.modality import Modality -from torch import nn - - -def _split_perturbations(config: BatchedPerturbationConfig, sizes: list[int]) -> list[BatchedPerturbationConfig]: - """Split a ``BatchedPerturbationConfig`` along the batch dimension.""" - it = iter(config.perturbations) - return [BatchedPerturbationConfig([next(it) for _ in range(s)]) for s in sizes] - - -def _merge_tensors(tensors: list[torch.Tensor | None]) -> torch.Tensor | None: - """Concatenate tensors along batch dim, or return None if all are None.""" - non_none = [t for t in tensors if t is not None] - if not non_none: - return None - return torch.cat(non_none, dim=0) - - -class BatchSplitAdapter(nn.Module): - """Wraps a model and splits batched forward calls into smaller chunks. - Has the same ``forward`` signature as ``X0Model``: - ``(video, audio, perturbations) -> (denoised_video, denoised_audio)``. - Args: - model: The model to wrap (``X0Model``, ``BlockStreamingWrapper``, etc.). - max_batch_size: Maximum batch size per forward pass. Input batches - larger than this are split into sequential chunks. - """ - - def __init__(self, model: nn.Module, max_batch_size: int) -> None: - if max_batch_size < 1: - raise ValueError(f"max_batch_size must be >= 1, got {max_batch_size}") - super().__init__() - self._model = model - self._max_batch_size = max_batch_size - - def _get_chunk_sizes(self, batch_size: int) -> list[int]: - full, remainder = divmod(batch_size, self._max_batch_size) - sizes = [self._max_batch_size] * full - if remainder: - sizes.append(remainder) - return sizes - - def forward( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - batch_size = (video or audio).latent.shape[0] - - if batch_size <= self._max_batch_size: - return self._model(video=video, audio=audio, perturbations=perturbations) - - sizes = self._get_chunk_sizes(batch_size) - n = len(sizes) - - v_chunks = video.split(sizes) if video is not None else [None] * n - a_chunks = audio.split(sizes) if audio is not None else [None] * n - p_chunks = _split_perturbations(perturbations, sizes) - - chunk_results = [self._model(video=vc, audio=ac, perturbations=pc) for vc, ac, pc in zip(v_chunks, a_chunks, p_chunks, strict=True)] - - results_v, results_a = zip(*chunk_results, strict=True) - return _merge_tensors(list(results_v)), _merge_tensors(list(results_a)) - - def __getattr__(self, name: str) -> Any: # noqa: ANN401 - """Proxy attribute access to the wrapped model.""" - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self._model, name) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/__init__.py deleted file mode 100644 index ee2d5c592..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -"""Block streaming: memory-efficient sequential-block inference. -Streams transformer blocks from safetensors to GPU one at a time. -Block weights are provided by a :class:`WeightsProvider` which handles -CPU-to-GPU copies, caching, and stream synchronization. Two weight -source strategies are available: -- **RAM streaming** (default): all blocks pre-loaded into pinned CPU - buffers with LoRA fusion at build time. Fast, higher CPU memory. -- **Disk streaming** (``cpu_slots < blocks_number``): blocks are read from - disk on demand by a :class:`DiskWeightSource`, on a background worker - thread. Slower, lower CPU memory. -""" - -from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder -from ltx_core.block_streaming.wrapper import BlockStreamingWrapper - -__all__ = [ - "DISK_CPU_SLOTS", - "BlockStreamingWrapper", - "StreamingModelBuilder", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/block_fetcher.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/block_fetcher.py deleted file mode 100644 index 96ab66841..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/block_fetcher.py +++ /dev/null @@ -1,82 +0,0 @@ -"""BlockFetcher: async disk reads on a worker thread.""" - -from __future__ import annotations - -import logging -import queue -import threading -from dataclasses import dataclass - -import torch -from ltx_core.block_streaming.disk import DiskBlockReader - -logger = logging.getLogger(__name__) - -Buffer = dict[str, torch.Tensor] - - -@dataclass(slots=True) -class FetchHandle: - """Caller-facing handle for an outstanding read returned by :meth:`BlockFetcher.submit`. - Carries only what the caller needs: a completion event the worker sets and the - read's error. The worker updates these once the read finishes. - """ - - _done: threading.Event - _error: BaseException | None = None - - def wait(self) -> BaseException | None: - """Block until the read finishes; return its error, or ``None`` on success.""" - self._done.wait() - return self._error - - -@dataclass(slots=True) -class _ReadRequest: - """One outstanding read, internal to :class:`BlockFetcher`. - The fetcher's worker reads block ``idx`` into the caller-carved ``buffer`` and - updates ``handle`` (its error, then its event) once the read has finished. - """ - - idx: int - buffer: Buffer - handle: FetchHandle - - -class BlockFetcher: - """Fills caller-supplied buffers on a worker thread.""" - - def __init__(self, reader: DiskBlockReader) -> None: - self._reader = reader - self._request_queue: queue.SimpleQueue[_ReadRequest | None] = queue.SimpleQueue() - self._worker = threading.Thread(target=self._run, name="BlockFetcher-IO", daemon=True) - self._worker.start() - - def submit(self, idx: int, buffer: Buffer) -> FetchHandle: - """Enqueue a read of block *idx* into the caller-carved *buffer*, return its handle.""" - handle = FetchHandle(_done=threading.Event()) - request = _ReadRequest(idx=idx, buffer=buffer, handle=handle) - self._request_queue.put(request) - return handle - - def cleanup(self) -> None: - """Drain pending reads, join the worker, close the reader.""" - self._request_queue.put(None) - self._worker.join() - self._reader.cleanup() - - def _run(self) -> None: - # Pinned buffers are allocated under the caller's inference_mode, so - # in-place copy_ from this thread requires inference_mode here too. - with torch.inference_mode(): - while True: - request = self._request_queue.get() - if request is None: - return - - try: - self._reader.read_into(request.buffer, request.idx) - except Exception as exc: - logger.exception("BlockFetcher: fetch failed for item %d", request.idx) - request.handle._error = exc - request.handle._done.set() diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/builder.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/builder.py deleted file mode 100644 index b8166deb7..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/builder.py +++ /dev/null @@ -1,463 +0,0 @@ -"""Builder that constructs a BlockStreamingWrapper from safetensors checkpoints.""" - -from __future__ import annotations - -import copy -import logging -from dataclasses import replace -from typing import TYPE_CHECKING, Final, Generic - -import safetensors -import torch -from ltx_core.block_streaming import utils as bs_utils -from ltx_core.block_streaming.block_fetcher import BlockFetcher -from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource -from ltx_core.block_streaming.pool import BufferPool -from ltx_core.block_streaming.provider import WeightsProvider -from ltx_core.block_streaming.source import DiskWeightSource, PinnedBlock, PinnedWeightSource, WeightSource -from ltx_core.block_streaming.utils import ( - carve_buffer, - derive_layout, - layout_nbytes, - make_block_key, - resolve_attr, -) -from ltx_core.block_streaming.wrapper import BlockStreamingWrapper -from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule, fuse_lora_weights -from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import ( - LoraPathStrengthAndSDOps, - LoraStateDictWithStrength, - ModelBuilderProtocol, - StateDict, - StateDictLoader, - TensorLayout, -) -from ltx_core.loader.registry import DummyRegistry, Registry -from ltx_core.loader.sd_ops import SDOps -from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader -from ltx_core.model.model_protocol import ModelConfigurator, ModelType -from torch import nn - -if TYPE_CHECKING: - from typing_extensions import Self - -logger = logging.getLogger(__name__) - -DISK_CPU_SLOTS = 2 -_DEFAULT_GPU_SLOTS = 2 -_PREFETCH_DEPTH = 2 - - -class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]): - """Immutable builder for :class:`BlockStreamingWrapper`. - Reads block weights from safetensors on demand. ``cpu_slots`` and - ``gpu_slots`` control the memory/speed trade-off (see :meth:`build`). - The builder is immutable (``with_*`` return modified copies) and exposes - its state via read-only properties backed by private attributes. - Args: - model_class_configurator: Creates the model from a config dict. - model_path: One or more ``.safetensors`` checkpoint paths. - model_sd_ops: Key remapping applied to safetensors keys. - module_ops: Module-level mutations for the meta model. - loras: LoRA adapters fused into weights at load time. - model_loader: Strategy for reading checkpoint metadata. - registry: Shared cache for loaded state dicts. - fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``; - use ``fp8_cast_fuse_rule`` for fp8_cast streaming so the pinned - buffers receive correctly-quantized weights. - blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g. - ``"transformer_blocks"``). - blocks_prefix: State-dict key prefix for block weights - (e.g. ``"transformer_blocks"``). - """ - - def __init__( - self, - model_class_configurator: type[ModelConfigurator[ModelType]], - model_path: str | tuple[str, ...], - model_sd_ops: SDOps | None = None, - module_ops: tuple[ModuleOps, ...] = (), - loras: tuple[LoraPathStrengthAndSDOps, ...] = (), - model_loader: StateDictLoader | None = None, - registry: Registry | None = None, - fuse_rule: FuseRule = bf16_fuse_rule, - blocks_attr: str = "", - blocks_prefix: str = "", - ) -> None: - # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute. - self._model_class_configurator: Final = model_class_configurator - self._model_path = model_path - self._model_sd_ops = model_sd_ops - self._module_ops = module_ops - self._loras = loras - self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader() - self._registry = registry if registry is not None else DummyRegistry() - self._fuse_rule = fuse_rule - self._blocks_attr = blocks_attr - self._blocks_prefix = blocks_prefix - - @property - def model_class_configurator(self) -> type[ModelConfigurator[ModelType]]: - return self._model_class_configurator - - @property - def model_path(self) -> str | tuple[str, ...]: - return self._model_path - - @property - def model_sd_ops(self) -> SDOps | None: - return self._model_sd_ops - - @property - def module_ops(self) -> tuple[ModuleOps, ...]: - return self._module_ops - - @property - def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]: - return self._loras - - @property - def model_loader(self) -> StateDictLoader: - return self._model_loader - - @property - def registry(self) -> Registry: - return self._registry - - @property - def fuse_rule(self) -> FuseRule: - return self._fuse_rule - - @property - def blocks_attr(self) -> str: - return self._blocks_attr - - @property - def blocks_prefix(self) -> str: - return self._blocks_prefix - - def with_sd_ops(self, sd_ops: SDOps | None) -> Self: - clone = copy.copy(self) - clone._model_sd_ops = sd_ops - return clone - - def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self: - clone = copy.copy(self) - clone._module_ops = module_ops - return clone - - def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self: - clone = copy.copy(self) - clone._loras = loras - return clone - - def with_registry(self, registry: Registry) -> Self: - clone = copy.copy(self) - clone._registry = registry - return clone - - def with_lora_load_device(self, device: torch.device) -> Self: - # Streaming fuses LoRAs into pinned CPU buffers; no other staging device is meaningful. - raise NotImplementedError("StreamingModelBuilder loads LoRA weights on CPU only.") - - def with_fuse_rule(self, fuse_rule: FuseRule) -> Self: - clone = copy.copy(self) - clone._fuse_rule = fuse_rule - return clone - - def model_config(self) -> dict: - """Read model configuration from the checkpoint metadata.""" - return read_model_config(self.model_path, self.model_loader) - - def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType: - """Create a model on the meta device and apply module operations.""" - return create_meta_model(self.model_class_configurator, config, module_ops) - - def build( - self, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - cpu_slots_count: int | None = None, - gpu_slots_count: int | None = None, - **_kwargs: object, - ) -> BlockStreamingWrapper: - """Build and return a ready-to-use :class:`BlockStreamingWrapper`. - Args: - device: GPU device for compute. ``None`` defaults to ``cuda``. - dtype: Weight dtype (e.g. ``torch.bfloat16``). Required. - cpu_slots_count: Number of pinned CPU buffer slots. - ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion). - gpu_slots_count: Number of GPU buffer slots. - ``None`` = ``_DEFAULT_GPU_SLOTS`` (2). - """ - if not self.blocks_prefix: - raise ValueError("blocks_prefix must be non-empty for streaming") - if dtype is None: - raise ValueError("StreamingModelBuilder.build requires an explicit dtype") - device = device if device is not None else torch.device("cuda") - - config = read_model_config(self.model_path, self.model_loader) - meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops) - meta_model.eval() - - blocks = resolve_attr(meta_model, self.blocks_attr) - - checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path] - block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix) - expected_indices = set(range(len(blocks))) - if set(block_key_map) != expected_indices: - missing = sorted(expected_indices - set(block_key_map)) - extra = sorted(set(block_key_map) - expected_indices) - raise ValueError(f"Block weights under prefix '{self.blocks_prefix}.' do not match the {len(blocks)} model blocks: missing indices {missing}, unexpected indices {extra}") - - cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks) - gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS - - if cpu_slots_count >= len(blocks): - lora_sd_and_strengths = self._load_lora_sds() - source, lora_sources = self._build_pinned_source(blocks, dtype, cpu_slots_count, block_key_map, lora_sd_and_strengths) - non_block_loras = lora_sd_and_strengths - else: - reader = DiskTensorReader(checkpoint_paths) - source, lora_sources = self._build_disk_source(blocks, dtype, cpu_slots_count, reader, block_key_map, prefetch_depth=_PREFETCH_DEPTH) - non_block_loras = [src.as_state_dict_with_strength() for src in lora_sources] - - self._load_non_block_weights(meta_model, non_block_keys, device, dtype, non_block_loras) - - copy_stream = torch.cuda.Stream(device=device) - gpu_pool = BufferPool( - source.slot_nbytes, - gpu_slots_count, - device, - reuse_barrier=lambda event: copy_stream.wait_event(event), - ) - provider = WeightsProvider( - gpu_pool, - copy_stream, - device, - source, - lora_sources, - self.blocks_prefix, - fuse_rule=self.fuse_rule, - ) - return BlockStreamingWrapper( - model=meta_model, - blocks=blocks, - provider=provider, - target_device=device, - ) - - def _load_lora_sds(self) -> list[LoraStateDictWithStrength]: - """Load each configured LoRA into a state dict for fusion (pinned path).""" - return [ - LoraStateDictWithStrength( - load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops), - lora.strength, - ) - for lora in self.loras - ] - - def _filtered_sd_ops(self, name_suffix: str, allowed_model_keys: frozenset[str]) -> SDOps: - """``model_sd_ops`` restricted to *allowed_model_keys* (post-rename keys). - The loader skips keys filtered to None before reading them, so a restricted - load never materializes the excluded partition. The distinct ``name`` avoids - a registry cache-id collision with the other partition. - """ - base = self.model_sd_ops if self.model_sd_ops is not None else SDOps("streaming").with_matching() - allowed = allowed_model_keys if base.allowed_keys is None else (allowed_model_keys & base.allowed_keys) - return replace(base, name=f"{base.name}__{name_suffix}", allowed_keys=allowed) - - def _build_pinned_source( - self, - blocks: nn.ModuleList, - dtype: torch.dtype, - cpu_slots_count: int, - block_key_map: dict[int, list[tuple[str, str]]], - lora_sd_and_strengths: list[LoraStateDictWithStrength], - ) -> tuple[WeightSource, list[LoraSource]]: - """Pre-load each block into its own contiguous pinned CPU buffer with LoRA fusion.""" - for block_idx in block_key_map: - if block_idx >= cpu_slots_count: - raise ValueError(f"Pinned source requires one CPU slot per block; got block index {block_idx} with only {cpu_slots_count} slots.") - - # One contiguous pinned buffer per block, carved into per-param views. The - # views (flattened by full key) are filled in place; the source then keeps - # only the contiguous buffer and the layout to re-carve it on read. - pinned_buffers: dict[int, torch.Tensor] = {} - block_layouts: dict[int, TensorLayout] = {} - fill_views: dict[str, torch.Tensor] = {} - for block_idx, entries in block_key_map.items(): - block_state = _block_state(blocks[block_idx]) - layout = derive_layout({param_name: block_state[param_name] for _sft_key, param_name in entries}, dtype) - buffer = bs_utils.alloc_buffer(layout_nbytes(layout), torch.device("cpu"), pin_memory=True) - views = carve_buffer(buffer, layout) - pinned_buffers[block_idx] = buffer - block_layouts[block_idx] = layout - for param_name, view in views.items(): - fill_views[make_block_key(self.blocks_prefix, block_idx, param_name)] = view - - block_sd = load_state_dict( - self.model_path, - self.model_loader, - self.registry, - torch.device("cpu"), - self._filtered_sd_ops("blocks", frozenset(fill_views)), - ) - - should_sync = False - for key, fused in fuse_lora_weights(block_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, preserve_input_device=False): - if key not in fill_views: - raise ValueError(f"Block-restricted load produced {key!r}, which is not a pinned block weight") - fill_views[key].copy_(fused, non_blocking=True) - block_sd.sd[key] = None - should_sync = True - if should_sync: - torch.cuda.synchronize() - - # Fill remaining pinned keys from the source state dict. - for key, view in fill_views.items(): - if block_sd.sd[key] is None: - continue - view.copy_(block_sd.sd[key]) - block_sd.sd[key] = None - - pinned = {idx: PinnedBlock(pinned_buffers[idx], block_layouts[idx]) for idx in pinned_buffers} - return PinnedWeightSource(pinned), [] - - def _build_disk_source( - self, - blocks: nn.ModuleList, - dtype: torch.dtype, - cpu_slots_count: int, - reader: DiskTensorReader, - block_key_map: dict[int, list[tuple[str, str]]], - prefetch_depth: int, - ) -> tuple[WeightSource, list[LoraSource]]: - """Create a DiskWeightSource backed by a DiskBlockReader. - Pool slots are sized to the largest block and carved per block on read, so - heterogeneous blocks (e.g. layers with differing attention layouts) share - one pool. Pool capacity is ``cpu_slots_count + prefetch_depth`` so the - lookahead loop in ``DiskWeightSource.get`` never evicts its own target. - Layouts come from the meta model; assumes module_ops keep the meta param - dtype in sync with the post-sd_ops checkpoint dtype. - """ - block_layouts = _block_layouts(blocks, block_key_map, dtype) - slot_nbytes = max(layout_nbytes(layout) for layout in block_layouts.values()) - - cpu_pool = BufferPool( - slot_nbytes, - cpu_slots_count + prefetch_depth, - torch.device("cpu"), - reuse_barrier=lambda event: event.synchronize(), - pin_memory=True, - ) - block_reader = DiskBlockReader( - reader=reader, - block_key_map=block_key_map, - sd_ops=self.model_sd_ops, - blocks_prefix=self.blocks_prefix, - ) - fetcher = BlockFetcher(block_reader) - source = DiskWeightSource( - cpu_pool, - fetcher, - block_layouts, - blocks_number=len(blocks), - prefetch_depth=prefetch_depth, - ) - lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras] - - return source, lora_sources - - @torch.inference_mode() - def _load_non_block_weights( - self, - model: nn.Module, - non_block_keys: list[tuple[str, str]], - device: torch.device, - dtype: torch.dtype, - lora_sd_and_strengths: list[LoraStateDictWithStrength], - ) -> None: - """Load the non-block weights onto *device* and fuse LoRAs -- both paths. - Reads through the loader with ``model_sd_ops`` restricted to the - non-block keys, so ``sd_ops`` (incl. kv-ops such as Gemma's ``lm_head`` - duplication) is applied exactly once and block tensors are never read. - """ - non_block_sd_ops = self._filtered_sd_ops("non_block", frozenset(mk for _sft_key, mk in non_block_keys)) - loaded = load_state_dict(self.model_path, self.model_loader, self.registry, device, non_block_sd_ops) - non_block_sd = {key: tensor.to(dtype=dtype) for key, tensor in loaded.sd.items()} - - if lora_sd_and_strengths: - non_block_state = StateDict(sd=non_block_sd, device=device, size=0, dtype={dtype}) - for key, fused in fuse_lora_weights( - non_block_state, - lora_sd_and_strengths, - fuse_rule=self.fuse_rule, - preserve_input_device=True, - ): - non_block_sd[key] = fused - - model.load_state_dict(non_block_sd, strict=False, assign=True) - - -def _block_state(block: nn.Module) -> dict[str, torch.Tensor]: - """Streamed-eligible tensors of a block: parameters then buffers. - Block streaming swaps both params and checkpoint-backed buffers (e.g. Gemma4's - per-layer ``layer_scalar``), so the layout, pinned packing, and meta-ordering - all consult parameters and buffers together. Non-checkpoint (computed) buffers - are harmless here -- only keys present in ``block_key_map`` are ever streamed. - """ - return {**dict(block.named_parameters()), **dict(block.named_buffers())} - - -def _block_layouts( - blocks: nn.ModuleList, - block_key_map: dict[int, list[tuple[str, str]]], - dtype: torch.dtype, -) -> dict[int, TensorLayout]: - """Per-block layout of the streamed tensors, taken from the meta model. - Blocks may differ in shape and even in which tensors they have (e.g. Gemma4's - full-attention layers drop ``v_proj``), so each block gets its own layout in - ``block_key_map`` order. The pinned packing, the disk reader, and the GPU carve - all key off this same per-block layout, so the provider's contiguous H2D copy - is valid for any entry order (no cross-block ordering required). - """ - layouts: dict[int, TensorLayout] = {} - for idx, entries in block_key_map.items(): - state = _block_state(blocks[idx]) - layouts[idx] = derive_layout({param_name: state[param_name] for _sft_key, param_name in entries}, dtype) - return layouts - - -def _scan_checkpoint_keys( - checkpoint_paths: list[str], - sd_ops: SDOps | None, - blocks_prefix: str, -) -> tuple[dict[int, list[tuple[str, str]]], list[tuple[str, str]]]: - """Partition checkpoint keys into per-block and non-block lists. - Opens the safetensors files for header-only key enumeration; no tensor data - is read. - """ - block_key_map: dict[int, list[tuple[str, str]]] = {} - non_block_keys: list[tuple[str, str]] = [] - prefix_dot = blocks_prefix + "." - for path in checkpoint_paths: - with safetensors.safe_open(path, framework="pt", device="cpu") as handle: - for sft_key in handle.keys(): # noqa: SIM118 - model_key = sd_ops.apply_to_key(sft_key) if sd_ops else sft_key - if model_key is None: - continue - if model_key.startswith(prefix_dot): - rest = model_key[len(prefix_dot) :] - idx_str, _, param_name = rest.partition(".") - try: - block_idx = int(idx_str) - except ValueError: - non_block_keys.append((sft_key, model_key)) - continue - block_key_map.setdefault(block_idx, []).append((sft_key, param_name)) - else: - non_block_keys.append((sft_key, model_key)) - return block_key_map, non_block_keys diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/disk.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/disk.py deleted file mode 100644 index 5ca3ccc54..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/disk.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Safetensors I/O and LoRA fusion for block streaming.""" - -from __future__ import annotations - -from collections.abc import Iterator - -import safetensors -import torch -from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key -from ltx_core.loader.fuse_loras import LoraProduct -from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict -from ltx_core.loader.sd_ops import SDOps - -_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = { - "F64": torch.float64, - "F32": torch.float32, - "F16": torch.float16, - "BF16": torch.bfloat16, -} - - -class DiskTensorReader: - """Key-based tensor accessor over one or more safetensors files.""" - - def __init__(self, paths: list[str]) -> None: - self._handles: list[safetensors.safe_open] = [] - self._key_to_handle_idx: dict[str, int] = {} - for path in paths: - handle = safetensors.safe_open(path, framework="pt", device="cpu") - handle_idx = len(self._handles) - self._handles.append(handle) - for sft_key in handle.keys(): # noqa: SIM118 - self._key_to_handle_idx[sft_key] = handle_idx - - def get_tensor(self, key: str) -> torch.Tensor: - return self._handles[self._key_to_handle_idx[key]].get_tensor(key) - - def close(self) -> None: - self._handles.clear() - self._key_to_handle_idx.clear() - - def __contains__(self, key: str) -> bool: - return key in self._key_to_handle_idx - - def __iter__(self) -> Iterator[str]: - return iter(self._key_to_handle_idx) - - -class DiskBlockReader: - """Reads one block at a time from safetensors into provided buffers.""" - - def __init__( - self, - reader: DiskTensorReader, - block_key_map: dict[int, list[tuple[str, str]]], - sd_ops: SDOps | None = None, - blocks_prefix: str = "", - ) -> None: - self._reader = reader - self._block_key_map = block_key_map - self._sd_ops = sd_ops - self._blocks_prefix = blocks_prefix - - def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None: - block_prefix = make_block_key(self._blocks_prefix, block_idx, "") - for sft_key, param_name in self._block_key_map[block_idx]: - tensor = self._reader.get_tensor(sft_key) - if self._sd_ops is None: - target[param_name].copy_(tensor) - continue - full_key = make_block_key(self._blocks_prefix, block_idx, param_name) - for result in self._sd_ops.apply_to_key_value(full_key, tensor): - if not result.new_key.startswith(block_prefix): - raise ValueError(f"SDOps output key '{result.new_key}' is outside block {block_idx} (expected prefix '{block_prefix}'); cannot route to a per-block buffer.") - target[result.new_key[len(block_prefix) :]].copy_(result.new_value) - - def cleanup(self) -> None: - self._reader.close() - - -class LoraSource: - """Pinned-memory cache of matched LoRA A/B factors backed by a single buffer.""" - - def __init__(self, path: str, sd_ops: SDOps | None, strength: float) -> None: - self.strength = strength - self._pinned_ab: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} - - a_keys: dict[str, str] = {} - b_keys: dict[str, str] = {} - with safetensors.safe_open(path, framework="pt", device="cpu") as handle: - for sft_key in handle.keys(): # noqa: SIM118 - model_key = sd_ops.apply_to_key(sft_key) if sd_ops is not None else sft_key - if model_key is None: - continue - if model_key.endswith(".lora_A.weight"): - a_keys[model_key[: -len(".lora_A.weight")]] = sft_key - elif model_key.endswith(".lora_B.weight"): - b_keys[model_key[: -len(".lora_B.weight")]] = sft_key - - matched_prefixes = list(a_keys.keys() & b_keys.keys()) - - # Build the layout from safetensors header metadata only — no tensor data is read. - layout: dict[str, tuple[torch.Size, torch.dtype]] = {} - for prefix in matched_prefixes: - a_slice_view = handle.get_slice(a_keys[prefix]) - b_slice_view = handle.get_slice(b_keys[prefix]) - layout[f"{prefix}.A"] = ( - torch.Size(a_slice_view.get_shape()), - _SAFETENSORS_DTYPE_TO_TORCH[a_slice_view.get_dtype()], - ) - layout[f"{prefix}.B"] = ( - torch.Size(b_slice_view.get_shape()), - _SAFETENSORS_DTYPE_TO_TORCH[b_slice_view.get_dtype()], - ) - - all_views = allocate_layout_views(layout, pin_memory=True) - - for prefix in matched_prefixes: - a_view = all_views[f"{prefix}.A"] - b_view = all_views[f"{prefix}.B"] - a_view.copy_(handle.get_tensor(a_keys[prefix])) - b_view.copy_(handle.get_tensor(b_keys[prefix])) - self._pinned_ab[prefix] = (a_view, b_view) - - def as_state_dict_with_strength(self) -> LoraStateDictWithStrength: - """Return a :class:`LoraStateDictWithStrength` view of the pinned A/B factors. - Lets non-block fusion consume the already-loaded disk-streaming LoRA - without re-reading the safetensors file or re-applying ``sd_ops``. - """ - sd: dict[str, torch.Tensor] = {} - for prefix, (a, b) in self._pinned_ab.items(): - sd[f"{prefix}.lora_A.weight"] = a - sd[f"{prefix}.lora_B.weight"] = b - size = sum(t.numel() * t.element_size() for t in sd.values()) - dtypes = {t.dtype for t in sd.values()} - return LoraStateDictWithStrength( - StateDict(sd=sd, device=torch.device("cpu"), size=size, dtype=dtypes), - self.strength, - ) - - def get_ab( - self, - param_prefix: str, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> LoraProduct | None: - """Return the :class:`LoraProduct` for *param_prefix*, or ``None``.""" - pair = self._pinned_ab.get(param_prefix) - if pair is None: - return None - a, b = pair - if device is not None and device.type == "cuda": - a = a.to(device=device, non_blocking=True) - b = b.to(device=device, non_blocking=True) - if dtype is not None: - a = a.to(dtype=dtype) - b = b.to(dtype=dtype) - return LoraProduct(a, b, self.strength) - - def cleanup(self) -> None: - self._pinned_ab.clear() diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/pool.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/pool.py deleted file mode 100644 index b242c8e8e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/pool.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Raw buffer pool for block streaming.""" - -from __future__ import annotations - -from collections import deque -from typing import Callable - -import torch -from ltx_core.block_streaming import utils - - -class BufferPool: - """Fixed pool of pre-allocated raw buffer slots with event-based reuse. - Slots are carved from a single contiguous ``uint8`` buffer; each is - ``slot_nbytes`` long and handed out as a raw 1-D ``uint8`` tensor. - Args: - slot_nbytes: Byte size of each slot. - capacity: Number of slots to pre-allocate. - device: Device for allocation. - reuse_barrier: Called with the pending event before a slot is reused. - pin_memory: Pin buffers (for async H2D copies from CPU). - """ - - def __init__( - self, - slot_nbytes: int, - capacity: int, - device: torch.device, - reuse_barrier: Callable[[torch.cuda.Event], None], - pin_memory: bool = False, - ) -> None: - self._slot_nbytes = slot_nbytes - self._capacity = capacity - self._free: deque[torch.Tensor] = deque() - self._events: dict[int, torch.cuda.Event] = {} - self._reuse_barrier = reuse_barrier - buffer = utils.alloc_buffer(max(slot_nbytes * capacity, 1), device, pin_memory) - for slot in range(capacity): - self._free.append(buffer[slot * slot_nbytes : (slot + 1) * slot_nbytes]) - - @property - def capacity(self) -> int: - return self._capacity - - @property - def slot_nbytes(self) -> int: - return self._slot_nbytes - - def acquire(self) -> torch.Tensor: - """Take a free raw slot, waiting any pending event before returning. - Raises :class:`RuntimeError` if every slot is currently in use. - """ - if not self._free: - raise RuntimeError(f"BufferPool exhausted: all {self._capacity} buffers are in use") - buffer = self._free.popleft() - event = self._events.pop(id(buffer), None) - if event is not None: - self._reuse_barrier(event) - return buffer - - def release(self, buffer: torch.Tensor, event: torch.cuda.Event | None = None) -> None: - """Return a raw slot to the free list. - The *buffer* must be the exact tensor object returned by :meth:`acquire` - (reuse is keyed on its identity). If *event* is given it is waited on the - next :meth:`acquire` of this slot, ensuring the prior operation finished. - """ - if event is not None: - self._events[id(buffer)] = event - self._free.append(buffer) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/provider.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/provider.py deleted file mode 100644 index 3cfd715a5..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/provider.py +++ /dev/null @@ -1,138 +0,0 @@ -"""GPU weights provider for block streaming.""" - -from __future__ import annotations - -from collections import OrderedDict -from typing import NamedTuple - -import torch -from ltx_core.block_streaming.disk import LoraSource -from ltx_core.block_streaming.pool import BufferPool -from ltx_core.block_streaming.source import WeightSource -from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes -from ltx_core.loader.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule -from ltx_core.loader.primitives import StateDict - -_EMPTY_STATE_DICT = StateDict(sd={}, device=torch.device("cpu"), size=0, dtype=set()) - - -class CachedBlock(NamedTuple): - """A cached GPU block: the raw pool slot plus the carved per-key views. - The raw slot is what is returned to the pool on eviction; the views are - what callers consume. - """ - - raw: torch.Tensor - views: dict[str, torch.Tensor] - - -class WeightsProvider: - """Provides GPU-ready block weights via H2D copy from a pinned CPU weight source. - Args: - pool: Pre-allocated GPU weight buffer pool. - copy_stream: Dedicated CUDA stream for async H2D copies. - target_device: GPU device for compute. - source: Pinned CPU weight source. - lora_sources: LoRA adapters fused on H2D copy. - blocks_prefix: State-dict prefix for LoRA key matching. - fuse_rule: Per-policy LoRA merge rule (must be streaming-compatible: - no companion-key emission). Defaults to ``bf16_fuse_rule``. - """ - - def __init__( - self, - pool: BufferPool, - copy_stream: torch.cuda.Stream, - target_device: torch.device, - source: WeightSource, - lora_sources: list[LoraSource] | None = None, - blocks_prefix: str = "", - fuse_rule: FuseRule = bf16_fuse_rule, - ) -> None: - self._copy_stream = copy_stream - self._pool = pool - self._cache: OrderedDict[int, CachedBlock] = OrderedDict() - self._events: dict[int, torch.cuda.Event] = {} - self._target_device = target_device - self._source = source - self._lora_sources = lora_sources or [] - self._blocks_prefix = blocks_prefix - self._fuse_rule = fuse_rule - - def get(self, idx: int) -> dict[str, torch.Tensor]: - """Return GPU weights for block *idx*. Does H2D copy on miss.""" - if idx in self._cache: - return self._cache[idx].views - - # Evict oldest GPU buffer if at capacity. - if len(self._cache) >= self._pool.capacity: - evicted_idx, evicted = self._cache.popitem(last=False) - self._pool.release(evicted.raw, event=self._events.pop(evicted_idx, None)) - - layout = self._source.block_layout(idx) - raw = self._pool.acquire() - gpu_weights = carve_buffer(raw, layout) - cpu_buffer = self._source.get(idx) - - h2d_event = self._copy_to_gpu(idx, raw, gpu_weights, cpu_buffer, layout_nbytes(layout)) - self._source.release(idx, event=h2d_event) - - self._cache[idx] = CachedBlock(raw, gpu_weights) - return gpu_weights - - def _copy_to_gpu( - self, - idx: int, - raw: torch.Tensor, - gpu_weights: dict[str, torch.Tensor], - cpu_buffer: torch.Tensor, - nbytes: int, - ) -> torch.cuda.Event: - """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute. - *cpu_buffer* is one contiguous source buffer carved by the same layout as - *raw*, so a single byte copy of its leading *nbytes* reproduces every view - in *gpu_weights*. The wait is intentionally inside this method so callers -- - and instrumentation regions wrapping it -- observe the full transfer time. - """ - if not cpu_buffer.is_contiguous() or cpu_buffer.dtype != torch.uint8 or cpu_buffer.numel() < nbytes: - raise ValueError(f"source buffer for block {idx} must be a contiguous uint8 buffer of >= {nbytes} bytes, got {cpu_buffer.dim()}-D {cpu_buffer.dtype} with {cpu_buffer.numel()} elements") - with torch.cuda.stream(self._copy_stream): - raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True) - if self._lora_sources: - self._fuse_block_loras(idx, gpu_weights) - h2d_event = torch.cuda.Event() - h2d_event.record(self._copy_stream) - - torch.cuda.current_stream(self._target_device).wait_event(h2d_event) - return h2d_event - - def release(self, idx: int, event: torch.cuda.Event) -> None: - """Attach a compute-done event -- waited before this buffer is recycled.""" - self._events[idx] = event - - def cleanup(self) -> None: - """Synchronize streams and release all resources.""" - self._copy_stream.synchronize() - torch.cuda.current_stream(self._target_device).synchronize() - self._cache.clear() - self._events.clear() - self._source.cleanup() - for lora in self._lora_sources: - lora.cleanup() - - def __len__(self) -> int: - return len(self._cache) - - def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None: - """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``.""" - agg_dtype = self._fuse_rule.aggregation_dtype - for name, tensor in weights.items(): - if not name.endswith(".weight"): - continue - prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight") - products = (ab for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources) if ab is not None) - deltas = aggregate_lora_products(products, agg_dtype) - if deltas is None: - continue - fused = self._fuse_rule(name, tensor, deltas, _EMPTY_STATE_DICT) - tensor.copy_(fused[name]) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/source.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/source.py deleted file mode 100644 index 63a33006a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/source.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Weight sources for block streaming: protocol and implementations.""" - -from __future__ import annotations - -from typing import NamedTuple, Protocol - -import torch -from ltx_core.block_streaming.block_fetcher import BlockFetcher, FetchHandle -from ltx_core.block_streaming.pool import BufferPool -from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes -from ltx_core.loader.primitives import TensorLayout - - -class WeightSource(Protocol): - """Provides pinned CPU weights for a given block index. - Blocks may be heterogeneous: each has its own layout, so the source exposes a - per-block layout and the byte size of the largest block (which sizes the - pool slots -- a smaller block is carved into the front of a max-sized slot). - The source is the single source of truth for each block's layout. - """ - - def block_layout(self, idx: int) -> TensorLayout: - """Per-block buffer layout (shape + dtype for each param).""" - ... - - @property - def slot_nbytes(self) -> int: - """Byte size of the largest block; sizes a pool slot (16-byte aligned).""" - ... - - def get(self, idx: int) -> torch.Tensor: - """Return one contiguous CPU buffer for block *idx*.""" - ... - - def release(self, idx: int, event: torch.cuda.Event | None) -> None: - """Signal that an async operation using these weights is guarded by *event*.""" - ... - - def cleanup(self) -> None: - """Release all resources (buffers, readers, events).""" - ... - - -class _Scheduled(NamedTuple): - """A scheduled (possibly in-flight) read: the raw pool slot and its fetch handle.""" - - raw: torch.Tensor - status: FetchHandle - - -class DiskWeightSource(WeightSource): - """WeightSource that streams blocks from disk via a :class:`BlockFetcher`. - ``get(idx)`` must be paired with a ``release(idx)`` before *idx* is fetched - again; getting a block that is already in flight raises. Each read acquires a - raw pool slot, carves it to block *idx*'s layout, and hands the carved views to - the fetcher to fill, so one max-sized slot can serve blocks of differing shapes. - ``get`` returns that same contiguous slot. - """ - - def __init__( - self, - pool: BufferPool, - fetcher: BlockFetcher, - block_layouts: dict[int, TensorLayout], - blocks_number: int, - prefetch_depth: int = 0, - ) -> None: - if blocks_number <= 0: - raise ValueError(f"blocks_number must be > 0, got {blocks_number}") - if prefetch_depth < 0: - raise ValueError(f"prefetch_depth must be >= 0, got {prefetch_depth}") - max_layout_nbytes = max((layout_nbytes(layout) for layout in block_layouts.values()), default=0) - if pool.slot_nbytes < max_layout_nbytes: - raise ValueError(f"pool slot is too small for the largest block: slot {pool.slot_nbytes} bytes < {max_layout_nbytes}") - - self._pool = pool - self._blocks_number = blocks_number - self._prefetch_depth = prefetch_depth - self._fetcher = fetcher - self._block_layouts = block_layouts - self._scheduled: dict[int, _Scheduled] = {} - self._in_flight: dict[int, torch.Tensor] = {} - - def block_layout(self, idx: int) -> TensorLayout: - return self._block_layouts[idx] - - @property - def slot_nbytes(self) -> int: - return self._pool.slot_nbytes - - def get(self, idx: int) -> torch.Tensor: - if idx in self._in_flight: - raise RuntimeError(f"Block {idx} is already in flight; release it before getting it again") - - scheduled = self._scheduled.pop(idx, None) - if scheduled is None: - scheduled = self._schedule(idx) - error = scheduled.status.wait() - if error is not None: - self._pool.release(scheduled.raw) - raise error - - self._in_flight[idx] = scheduled.raw - for k in range(1, self._prefetch_depth + 1): - self._ensure_scheduled((idx + k) % self._blocks_number) - return scheduled.raw - - def release(self, idx: int, event: torch.cuda.Event | None) -> None: - raw_buffer = self._in_flight.pop(idx) - self._pool.release(raw_buffer, event=event) - - def cleanup(self) -> None: - self._fetcher.cleanup() - while self._in_flight: - _, raw_buffer = self._in_flight.popitem() - self._pool.release(raw_buffer) - while self._scheduled: - _, scheduled = self._scheduled.popitem() - scheduled.status.wait() - self._pool.release(scheduled.raw) - - def _ensure_scheduled(self, idx: int) -> None: - """Schedule a read for *idx* if one is not already pending.""" - if idx not in self._scheduled: - self._scheduled[idx] = self._schedule(idx) - - def _schedule(self, idx: int) -> _Scheduled: - """Acquire a raw slot, carve it to block *idx*, enqueue a read, return the handle. - The raw slot and its fetch status are returned together so the caller can - track both as one unit; the fetcher only receives the carved views to fill. - """ - raw_buffer = self._pool.acquire() - carved = carve_buffer(raw_buffer, self._block_layouts[idx]) - status = self._fetcher.submit(idx, carved) - return _Scheduled(raw_buffer, status) - - -class PinnedBlock(NamedTuple): - """A pre-loaded pinned block: its single contiguous buffer and the layout to carve it with.""" - - buffer: torch.Tensor - layout: TensorLayout - - -class PinnedWeightSource(WeightSource): - """Pre-loaded pinned CPU weights, one contiguous (possibly heterogeneous) buffer per block.""" - - def __init__(self, blocks: dict[int, PinnedBlock]) -> None: - if not blocks: - raise ValueError("PinnedWeightSource requires at least one block") - self._blocks = blocks - self._slot_nbytes = max(layout_nbytes(block.layout) for block in blocks.values()) - - def block_layout(self, idx: int) -> TensorLayout: - return self._blocks[idx].layout - - @property - def slot_nbytes(self) -> int: - return self._slot_nbytes - - def get(self, idx: int) -> torch.Tensor: - return self._blocks[idx].buffer - - def release(self, idx: int, event: torch.cuda.Event | None) -> None: - pass - - def cleanup(self) -> None: - self._blocks.clear() - - def __len__(self) -> int: - return len(self._blocks) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/utils.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/utils.py deleted file mode 100644 index 5b72aa15d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/utils.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Shared utilities for the block_streaming package.""" - -from __future__ import annotations - -import math -import weakref -from dataclasses import dataclass -from typing import Any, NamedTuple - -import torch -from ltx_core.loader.primitives import TensorLayout -from torch import nn - -FP8_DTYPES = frozenset({torch.float8_e4m3fn, torch.float8_e5m2}) - -_BUFFER_ALIGN = 16 - - -def make_block_key(blocks_prefix: str, block_idx: int, param_name: str) -> str: - """Return the state-dict key for *param_name* under block *block_idx*.""" - return f"{blocks_prefix}.{block_idx}.{param_name}" - - -def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList: - """Resolve a dotted attribute path like ``'model.language_model.layers'``.""" - obj: Any = module - for part in dotted_path.split("."): - obj = getattr(obj, part) - if not isinstance(obj, nn.ModuleList): - raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}") - return obj - - -def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Tensor) -> None: - """Assign *tensor* to the parameter/buffer at *dotted_name* inside *root*. - Unlike ``param.data = tensor``, this works even when the existing parameter - lives on the ``meta`` device (which has an incompatible storage type). - """ - parts = dotted_name.split(".") - parent = root - for part in parts[:-1]: - parent = getattr(parent, part) - leaf = parts[-1] - if leaf in parent._parameters: - parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=False) - elif leaf in parent._buffers: - parent._buffers[leaf] = tensor - else: - raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}") - - -def derive_layout(tensors: dict[str, torch.Tensor], dtype: torch.dtype | None = None) -> TensorLayout: - """Derive a layout from a ``{name: tensor}`` dict. - If ``dtype`` is given, non-FP8 dtypes are coerced to it (FP8 preserved). If - ``None``, the source dtype is preserved as-is. - """ - return {name: (t.shape, t.dtype if dtype is None or t.dtype in FP8_DTYPES else dtype) for name, t in tensors.items()} - - -def _align_up(offset: int, alignment: int) -> int: - return (offset + alignment - 1) & ~(alignment - 1) - - -def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None: - """Allocate exactly ``nbytes`` of pinned host memory via ``cudaHostRegister``. - Bypasses PyTorch's ``CachingHostAllocator``, which rounds every - ``pin_memory=True`` request up to ``PowerOf2Ceil(N)`` (see - ``aten/src/ATen/core/CachingHostAllocator.h``). Returns ``None`` if - registration fails. The unregister hook is bound to the storage (not the - tensor) so views of the buffer keep the registration alive until the - memory is actually freed. Caller is responsible for ensuring CUDA is - available. - """ - cudart = torch.cuda.cudart() - buf = torch.empty(nbytes, dtype=torch.uint8) - ptr = buf.data_ptr() - err = int(cudart.cudaHostRegister(ptr, nbytes, 0)) - if err != 0: - return None - weakref.finalize(buf.untyped_storage(), lambda p=ptr: cudart.cudaHostUnregister(p)) - return buf - - -def alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor: - """Allocate one ``uint8`` buffer for :func:`allocate_layout_views`. - For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching - allocator's power-of-2 rounding. Falls back to the caching allocator if - registration fails. Raises if pinning is requested without a CUDA runtime, - since pinning is fundamentally a CUDA driver operation. - """ - if pin_memory and (device is None or torch.device(device).type == "cpu"): - if not torch.cuda.is_available(): - raise RuntimeError("pin_memory=True requires CUDA, which is not available") - buf = _alloc_pinned_exact(nbytes) - if buf is not None: - return buf - return torch.empty(nbytes, dtype=torch.uint8, device=device, pin_memory=pin_memory) - - -@dataclass(frozen=True) -class _TensorSlice: - """Location of a single tensor view within the buffer.""" - - offset: int - shape: torch.Size - dtype: torch.dtype - - def size(self) -> int: - return math.prod(self.shape) * self.dtype.itemsize - - -class LayoutSlices(NamedTuple): - """Per-key tensor slices of a layout plus the total aligned buffer size.""" - - slices: dict[str, _TensorSlice] - nbytes: int - - -def _layout_slices(layout: TensorLayout) -> LayoutSlices: - """Compute the byte offset of each key in *layout* and the total aligned size. - The size is at least one byte so empty layouts still produce a valid buffer. - """ - slices: dict[str, _TensorSlice] = {} - cursor = 0 - for key, (shape, dtype) in layout.items(): - cursor = _align_up(cursor, _BUFFER_ALIGN) - slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype) - cursor += slices[key].size() - return LayoutSlices(slices, max(_align_up(cursor, _BUFFER_ALIGN), 1)) - - -def layout_nbytes(layout: TensorLayout) -> int: - """Byte size of one contiguous, 16-byte-aligned buffer holding *layout* (>= 1).""" - return _layout_slices(layout).nbytes - - -def carve_buffer(buffer: torch.Tensor, layout: TensorLayout) -> dict[str, torch.Tensor]: - """Carve per-key tensor views for *layout* into the front of *buffer*. - *buffer* is a 1-D ``uint8`` tensor at least :func:`layout_nbytes` long. Each - returned tensor is a non-overlapping slice of its leading bytes reinterpreted - at the requested shape and dtype; any trailing bytes are left unused. That - slack is what lets one max-sized pool slot hold a smaller (heterogeneous) - block. The views keep *buffer*'s storage alive via PyTorch refcounting. - """ - if buffer.dtype != torch.uint8 or buffer.dim() != 1: - raise ValueError(f"carve_buffer expects a 1-D uint8 buffer, got {buffer.dim()}-D {buffer.dtype}") - slices, nbytes = _layout_slices(layout) - if buffer.numel() < nbytes: - raise ValueError(f"buffer too small to carve layout: need {nbytes} bytes, got {buffer.numel()}") - return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()} - - -def allocate_layout_views( - layout: TensorLayout, - device: torch.device | None = None, - pin_memory: bool = False, -) -> dict[str, torch.Tensor]: - """Allocate a single ``uint8`` buffer and return per-key tensor views into it. - All keys in *layout* live in one contiguous allocation; each returned - tensor is a non-overlapping slice of that buffer reinterpreted at the - requested shape and dtype. The views keep the underlying storage alive - via PyTorch refcounting — drop them all to release the memory. - """ - buffer = alloc_buffer(layout_nbytes(layout), device, pin_memory) - return carve_buffer(buffer, layout) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/wrapper.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/wrapper.py deleted file mode 100644 index 148e2928f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/wrapper.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Block streaming wrapper: streams transformer blocks through a WeightsProvider.""" - -from __future__ import annotations - -import itertools -from typing import Any - -import torch -from ltx_core.block_streaming.provider import WeightsProvider -from ltx_core.block_streaming.utils import assign_tensor_to_module -from torch import nn - - -class BlockStreamingWrapper(nn.Module): - """Streams sequential model blocks through GPU buffer caches. - The wrapper delegates all weight management to a :class:`WeightsProvider` - which handles CPU-to-GPU copies, caching, LoRA fusion, and stream - synchronization internally. - Use :class:`StreamingModelBuilder` to construct this wrapper -- it - handles checkpoint parsing, source selection, and provider creation. - Args: - model: The wrapped model (non-block params already on GPU). - blocks: Sequential blocks to stream (``nn.ModuleList``). - provider: Provides GPU-ready weights on demand. - target_device: GPU device for compute. - """ - - def __init__( - self, - model: nn.Module, - blocks: nn.ModuleList, - provider: WeightsProvider, - target_device: torch.device, - ) -> None: - super().__init__() - self._model = model - self._blocks = blocks - self._target_device = target_device - self._provider = provider - - self._hooks: list[torch.utils.hooks.RemovableHandle] = [] - self._register_hooks() - - # ------------------------------------------------------------------ - # Hook registration - # ------------------------------------------------------------------ - - def _pre_hook(self, block_idx: int) -> None: - """Load GPU weights for a block and inject them into its parameters.""" - gpu_weights = self._provider.get(block_idx) - - block = self._blocks[block_idx] - for name, _param in itertools.chain(block.named_parameters(), block.named_buffers()): - assign_tensor_to_module(block, name, gpu_weights[name]) - - def _post_hook(self, block_idx: int) -> None: - """Record a compute-done event and release the block weights.""" - compute_done = torch.cuda.Event() - compute_done.record(torch.cuda.current_stream(self._target_device)) - self._provider.release(block_idx, event=compute_done) - - def _register_hooks(self) -> None: - for idx, block in enumerate(self._blocks): - pre = block.register_forward_pre_hook( - lambda _mod, _args, *, idx=idx: self._pre_hook(idx), - ) - post = block.register_forward_hook( - lambda _mod, _args, _out, *, idx=idx: self._post_hook(idx), - ) - self._hooks.extend([pre, post]) - - # ------------------------------------------------------------------ - # Teardown - # ------------------------------------------------------------------ - - def teardown(self) -> None: - """Remove hooks and release all resources.""" - for h in self._hooks: - h.remove() - self._hooks.clear() - self._provider.cleanup() - - # ------------------------------------------------------------------ - # Forward and attribute delegation - # ------------------------------------------------------------------ - - def forward(self, *args: Any, **kwargs: Any) -> Any: # noqa: ANN401 - return self._model(*args, **kwargs) - - def __getattr__(self, name: str) -> Any: # noqa: ANN401 - """Proxy attribute access to the wrapped model.""" - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self._model, name) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/__init__.py deleted file mode 100644 index c1cb638b4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -""" -Diffusion pipeline components. -Submodules: - diffusion_steps - Diffusion stepping algorithms (EulerDiffusionStep) - guiders - Guidance strategies (CFGGuider, STGGuider, APG variants) - noisers - Noise samplers (GaussianNoiser) - patchifiers - Latent patchification (VideoLatentPatchifier, AudioPatchifier) - protocols - Protocol definitions (Patchifier, etc.) - schedulers - Sigma schedulers (LTX2Scheduler, LinearQuadraticScheduler) -""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/diffusion_steps.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/diffusion_steps.py deleted file mode 100644 index a0ebc5f6b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/diffusion_steps.py +++ /dev/null @@ -1,183 +0,0 @@ -import torch -from ltx_core.components.protocols import DiffusionStepProtocol -from ltx_core.utils import to_velocity - - -def _get_ancestral_step( - sigma_from: torch.Tensor, - sigma_to: torch.Tensor, - eta: float = 1.0, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute ``(sigma_down, sigma_up)`` for one DDIM ancestral sampling step. - Both inputs are in the rescaled parameterization ``sigma / alpha``. - Returns ``sigma_down`` (deterministic component) and ``sigma_up`` - (stochastic component) in the same rescaled space. - """ - if not eta: - return sigma_to, torch.zeros_like(sigma_to) - variance = sigma_to**2 * (sigma_from**2 - sigma_to**2).clamp(min=0) / sigma_from**2 - sigma_up = (eta * variance**0.5).clamp(max=sigma_to) - sigma_down = (sigma_to**2 - sigma_up**2).clamp(min=0) ** 0.5 - return sigma_down, sigma_up - - -class EulerDiffusionStep(DiffusionStepProtocol): - """ - First-order Euler method for diffusion sampling. - Takes a single step from the current noise level (sigma) to the next by - computing velocity from the denoised prediction and applying: sample + velocity * dt. - """ - - def step(self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **_kwargs) -> torch.Tensor: - sigma = sigmas[step_index] - sigma_next = sigmas[step_index + 1] - dt = sigma_next - sigma - velocity = to_velocity(sample, sigma, denoised_sample) - - return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype) - - -class Res2sDiffusionStep(DiffusionStepProtocol): - """ - Second-order diffusion step for res_2s sampling with SDE noise injection. - Used by the res_2s denoising loop. Advances the sample from the current - sigma to the next by mixing a deterministic update (from the denoised - prediction) with injected noise via ``get_sde_coeff``, producing - variance-preserving transitions. - """ - - @staticmethod - def get_sde_coeff( - sigma_next: torch.Tensor, - sigma_up: torch.Tensor | None = None, - sigma_down: torch.Tensor | None = None, - sigma_max: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Compute SDE coefficients (alpha_ratio, sigma_down, sigma_up) for the step. - Given either ``sigma_down`` or ``sigma_up``, returns the mixing - coefficients used for variance-preserving noise injection. If - ``sigma_up`` is provided, ``sigma_down`` and ``alpha_ratio`` are - derived; if ``sigma_down`` is provided, ``sigma_up`` and - ``alpha_ratio`` are derived. - """ - if sigma_down is not None: - alpha_ratio = (1 - sigma_next) / (1 - sigma_down) - sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(min=0) ** 0.5 - elif sigma_up is not None: - # Fallback to avoid sqrt(neg_num) - sigma_up.clamp_(max=sigma_next * 0.9999) - sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next) - sigma_signal = sigmax - sigma_next - sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5 - alpha_ratio = sigma_signal + sigma_residual - sigma_down = sigma_residual / alpha_ratio - else: - alpha_ratio = torch.ones_like(sigma_next) - sigma_down = sigma_next - sigma_up = torch.zeros_like(sigma_next) - - sigma_up = torch.nan_to_num(sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0) - # Replace NaNs in sigma_down with corresponding sigma_next elements (float32) - nan_mask = torch.isnan(sigma_down) - sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype) - alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0) - - return alpha_ratio, sigma_down, sigma_up - - def step( - self, - sample: torch.Tensor, - denoised_sample: torch.Tensor, - sigmas: torch.Tensor, - step_index: int, - noise: torch.Tensor, - eta: float = 0.5, - ) -> torch.Tensor: - """Advance one step with SDE noise injection via get_sde_coeff. - Args: - sample: Current noisy sample. - denoised_sample: Denoised prediction from the model. - sigmas: Noise schedule tensor. - step_index: Current step index in the schedule. - noise: Random noise tensor for stochastic injection. - eta: Controls stochastic noise injection strength (0=deterministic, 1=maximum). Default 0.5. - Returns: - Next sample with SDE noise injection applied. - """ - sigma = sigmas[step_index] - sigma_next = sigmas[step_index + 1] - alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * eta) - output_dtype = denoised_sample.dtype - if torch.any(sigma_up == 0) or torch.any(sigma_next == 0): - return denoised_sample - - # Extract epsilon prediction - eps_next = (sample - denoised_sample) / (sigma - sigma_next) - denoised_next = sample - sigma * eps_next - - # Mix deterministic and stochastic components - x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise - return x_noised.to(output_dtype) - - -class EulerCfgPpDiffusionStep(DiffusionStepProtocol): - """Euler step using the CFG++ correction for the ODE derivative. - Instead of the standard velocity formula, the ODE derivative is computed - from the unconditioned prediction, keeping the conditioned prediction as - the target denoised state. Ancestral (DDIM) noise injection is applied - in the rescaled sigma parameterization (sigma / alpha). - All diffusion quantities (alpha, ODE derivative, ancestral coefficients) - are computed internally from ``sigmas`` and ``uncond_denoised``. - Reference: CFG++ (https://arxiv.org/abs/2406.08070). - """ - - def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None: - self.eta = eta - self.s_noise = s_noise - - def step( - self, - sample: torch.Tensor, - denoised_sample: torch.Tensor, - sigmas: torch.Tensor, - step_index: int, - uncond_denoised: torch.Tensor, - noise: torch.Tensor | None = None, - **_kwargs, - ) -> torch.Tensor: - """Advance one CFG++ Euler step. - Args: - sample: Current noisy latent x_t. - denoised_sample: Conditioned denoised prediction x_0^cond. - sigmas: Full sigma schedule tensor. - step_index: Current step index. - uncond_denoised: Unconditioned denoised prediction x_0^uncond, - used to compute the ODE derivative direction. - noise: Noise tensor for stochastic injection; ignored when - ``eta=0`` or ``s_noise=0``. - Returns: - Updated latent x_{t-1}. - """ - sigma_s = sigmas[step_index].to(torch.float32) - sigma_t = sigmas[step_index + 1].to(torch.float32) - _eps = torch.finfo(torch.float32).eps - # Clamp to avoid division by zero when sigma == 1.0 exactly. - alpha_s = (1.0 - sigma_s).clamp(min=_eps) - alpha_t = (1.0 - sigma_t).clamp(min=_eps) - - x = sample.to(torch.float32) - denoised = denoised_sample.to(torch.float32) - uncond = uncond_denoised.to(torch.float32) - - # ODE derivative: direction toward noise using uncond prediction (CFG++ correction) - d = (x - alpha_s * uncond) / sigma_s - - # Ancestral step in rescaled sigma space (sigma / alpha) - sigma_down, sigma_up = _get_ancestral_step(sigma_s / alpha_s, sigma_t / alpha_t, eta=self.eta) - sigma_down = alpha_t * sigma_down - - x_next = alpha_t * denoised + sigma_down * d - if noise is not None and self.eta > 0 and self.s_noise > 0: - x_next = x_next + alpha_t * noise.to(torch.float32) * self.s_noise * sigma_up - return x_next.to(sample.dtype) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/guiders.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/guiders.py deleted file mode 100644 index 53765b779..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/guiders.py +++ /dev/null @@ -1,361 +0,0 @@ -import math -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field - -import torch -from ltx_core.components.protocols import GuiderProtocol - - -@dataclass(frozen=True) -class CFGGuider(GuiderProtocol): - """ - Classifier-free guidance (CFG) guider. - Computes the guidance delta as (scale - 1) * (cond - uncond), steering the - denoising process toward the conditioned prediction. - Attributes: - scale: Guidance strength. 1.0 means no guidance, higher values increase - adherence to the conditioning. - """ - - scale: float - - def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: - return (self.scale - 1) * (cond - uncond) - - def enabled(self) -> bool: - return self.scale != 1.0 - - -@dataclass(frozen=True) -class CFGStarRescalingGuider(GuiderProtocol): - """ - Calculates the CFG delta between conditioned and unconditioned samples. - To minimize offset in the denoising direction and move mostly along the - conditioning axis within the distribution, the unconditioned sample is - rescaled in accordance with the norm of the conditioned sample. - Attributes: - scale (float): - Global guidance strength. A value of 1.0 corresponds to no extra - guidance beyond the base model prediction. Values > 1.0 increase - the influence of the conditioned sample relative to the - unconditioned one. - """ - - scale: float - - def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: - rescaled_neg = projection_coef(cond, uncond) * uncond - return (self.scale - 1) * (cond - rescaled_neg) - - def enabled(self) -> bool: - return self.scale != 1.0 - - -@dataclass(frozen=True) -class STGGuider(GuiderProtocol): - """ - Calculates the STG delta between conditioned and perturbed denoised samples. - Perturbed samples are the result of the denoising process with perturbations, - e.g. attentions acting as passthrough for certain layers and modalities. - Attributes: - scale (float): - Global strength of the STG guidance. A value of 0.0 disables the - guidance. Larger values increase the correction applied in the - direction of (pos_denoised - perturbed_denoised). - """ - - scale: float - - def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor: - return self.scale * (pos_denoised - perturbed_denoised) - - def enabled(self) -> bool: - return self.scale != 0.0 - - -@dataclass(frozen=True) -class LtxAPGGuider(GuiderProtocol): - """ - Calculates the APG (adaptive projected guidance) delta between conditioned - and unconditioned samples. - To minimize offset in the denoising direction and move mostly along the - conditioning axis within the distribution, the (cond - uncond) delta is - decomposed into components parallel and orthogonal to the conditioned - sample. The `eta` parameter weights the parallel component, while `scale` - is applied to the orthogonal component. Optionally, a norm threshold can - be used to suppress guidance when the magnitude of the correction is small. - Attributes: - scale (float): - Strength applied to the component of the guidance that is orthogonal - to the conditioned sample. Controls how aggressively we move in - directions that change semantics but stay consistent with the - conditioning manifold. - eta (float): - Weight of the component of the guidance that is parallel to the - conditioned sample. A value of 1.0 keeps the full parallel - component; values in [0, 1] attenuate it, and values > 1.0 amplify - motion along the conditioning direction. - norm_threshold (float): - Minimum L2 norm of the guidance delta below which the guidance - can be reduced or ignored (depending on implementation). - This is useful for avoiding noisy or unstable updates when the - guidance signal is very small. - """ - - scale: float - eta: float = 1.0 - norm_threshold: float = 0.0 - - def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: - guidance = cond - uncond - if self.norm_threshold > 0: - ones = torch.ones_like(guidance) - guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True) - scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm) - guidance = guidance * scale_factor - proj_coeff = projection_coef(guidance, cond) - g_parallel = proj_coeff * cond - g_orth = guidance - g_parallel - g_apg = g_parallel * self.eta + g_orth - - return g_apg * (self.scale - 1) - - def enabled(self) -> bool: - return self.scale != 1.0 - - -@dataclass(frozen=False) -class LegacyStatefulAPGGuider(GuiderProtocol): - """ - Calculates the APG (adaptive projected guidance) delta between conditioned - and unconditioned samples. - To minimize offset in the denoising direction and move mostly along the - conditioning axis within the distribution, the (cond - uncond) delta is - decomposed into components parallel and orthogonal to the conditioned - sample. The `eta` parameter weights the parallel component, while `scale` - is applied to the orthogonal component. Optionally, a norm threshold can - be used to suppress guidance when the magnitude of the correction is small. - Attributes: - scale (float): - Strength applied to the component of the guidance that is orthogonal - to the conditioned sample. Controls how aggressively we move in - directions that change semantics but stay consistent with the - conditioning manifold. - eta (float): - Weight of the component of the guidance that is parallel to the - conditioned sample. A value of 1.0 keeps the full parallel - component; values in [0, 1] attenuate it, and values > 1.0 amplify - motion along the conditioning direction. - norm_threshold (float): - Minimum L2 norm of the guidance delta below which the guidance - can be reduced or ignored (depending on implementation). - This is useful for avoiding noisy or unstable updates when the - guidance signal is very small. - momentum (float): - Exponential moving-average coefficient for accumulating guidance - over time. running_avg = momentum * running_avg + guidance - """ - - scale: float - eta: float - norm_threshold: float = 5.0 - momentum: float = 0.0 - # it is user's responsibility not to use same APGGuider for several denoisings or different modalities - # in order not to share accumulated average across different denoisings or modalities - running_avg: torch.Tensor | None = None - - def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: - guidance = cond - uncond - if self.momentum != 0: - if self.running_avg is None: - self.running_avg = guidance.clone() - else: - self.running_avg = self.momentum * self.running_avg + guidance - guidance = self.running_avg - - if self.norm_threshold > 0: - ones = torch.ones_like(guidance) - guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True) - scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm) - guidance = guidance * scale_factor - - proj_coeff = projection_coef(guidance, cond) - g_parallel = proj_coeff * cond - g_orth = guidance - g_parallel - g_apg = g_parallel * self.eta + g_orth - - return g_apg * self.scale - - def enabled(self) -> bool: - return self.scale != 0.0 - - -@dataclass(frozen=True) -class MultiModalGuiderParams: - """ - Parameters for the multi-modal guider. - """ - - cfg_scale: float = 1.0 - "CFG (Classifier-free guidance) scale controlling how strongly the model adheres to the prompt." - stg_scale: float = 0.0 - "STG (Spatio-Temporal Guidance) scale controls how strongly the model reacts to the perturbation of the modality." - stg_blocks: list[int] | None = field(default_factory=list) - "Which transformer blocks to perturb for STG." - rescale_scale: float = 0.0 - "Rescale scale controlling how strongly the model rescales the modality after applying other guidance." - modality_scale: float = 1.0 - "Modality scale controlling how strongly the model reacts to the perturbation of the modality." - skip_step: int = 0 - "Skip step controlling how often the model skips the step." - - -def _params_for_sigma_from_sorted_dict(sigma: float, params_by_sigma: Sequence[tuple[float, MultiModalGuiderParams]]) -> MultiModalGuiderParams: - """ - Return params for the given sigma from a sorted (sigma_upper_bound -> params) structure. - Keys are sorted descending (bin upper bounds). Bin i is (key_{i+1}, key_i]. - Get all keys >= sigma; use last in list (smallest such key = upper bound of bin containing sigma), - or last entry in the sequence if list is empty (sigma above max key). - """ - if not params_by_sigma: - raise ValueError("params_by_sigma must be non-empty") - sigma = float(sigma) - keys_desc = [k for k, _ in params_by_sigma] - keys_ge_sigma = [k for k in keys_desc if k >= sigma] - # sigma above all keys: use first bin (max key) - key = keys_ge_sigma[-1] if keys_ge_sigma else keys_desc[0] - return next(p for k, p in params_by_sigma if k == key) - - -@dataclass(frozen=True) -class MultiModalGuider: - """ - Multi-modal guider with constant params per instance. - For sigma-dependent params, use MultiModalGuiderFactory.build_from_sigma(sigma) to - obtain a guider for each step. - """ - - params: MultiModalGuiderParams - negative_context: torch.Tensor | None = None - - def calculate( - self, - cond: torch.Tensor, - uncond_text: torch.Tensor | float, - uncond_perturbed: torch.Tensor | float, - uncond_modality: torch.Tensor | float, - ) -> torch.Tensor: - """ - The guider calculates the guidance delta as (scale - 1) * (cond - uncond) for cfg and modality cfg, - and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned - prediction. - """ - dtype = cond.dtype - cond = cond.float() - uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text - uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed - uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality - pred = cond + (self.params.cfg_scale - 1) * (cond - uncond_text) + self.params.stg_scale * (cond - uncond_perturbed) + (self.params.modality_scale - 1) * (cond - uncond_modality) - - if self.params.rescale_scale != 0: - factor = cond.std() / pred.std() - factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale) - pred = pred * factor - - return pred.to(dtype) - - def do_unconditional_generation(self) -> bool: - """Returns True if the guider is doing unconditional generation.""" - return not math.isclose(self.params.cfg_scale, 1.0) - - def do_perturbed_generation(self) -> bool: - """Returns True if the guider is doing perturbed generation.""" - return not math.isclose(self.params.stg_scale, 0.0) - - def do_isolated_modality_generation(self) -> bool: - """Returns True if the guider is doing isolated modality generation.""" - return not math.isclose(self.params.modality_scale, 1.0) - - def should_skip_step(self, step: int) -> bool: - """Returns True if the guider should skip the step.""" - if self.params.skip_step == 0: - return False - return step % (self.params.skip_step + 1) != 0 - - -@dataclass(frozen=True) -class MultiModalGuiderFactory: - """ - Factory that creates a MultiModalGuider for a given sigma. - Single source of truth: _params_by_sigma (schedule). Use constant() for - one params for all sigma, from_dict() for sigma-binned params. - """ - - negative_context: torch.Tensor | None = None - _params_by_sigma: tuple[tuple[float, MultiModalGuiderParams], ...] = () - - @classmethod - def constant( - cls, - params: MultiModalGuiderParams, - negative_context: torch.Tensor | None = None, - ) -> "MultiModalGuiderFactory": - """Build a factory with constant params (same guider for all sigma).""" - return cls( - negative_context=negative_context, - _params_by_sigma=((float("inf"), params),), - ) - - @classmethod - def from_dict( - cls, - sigma_to_params: Mapping[float, MultiModalGuiderParams], - negative_context: torch.Tensor | None = None, - ) -> "MultiModalGuiderFactory": - """ - Build a factory from a dict of sigma_value -> MultiModalGuiderParams. - Keys are sorted descending and used for bin lookup in params(sigma). - """ - if not sigma_to_params: - raise ValueError("sigma_to_params must be non-empty") - sorted_items = tuple(sorted(sigma_to_params.items(), key=lambda x: x[0], reverse=True)) - return cls(negative_context=negative_context, _params_by_sigma=sorted_items) - - def params(self, sigma: float | torch.Tensor) -> MultiModalGuiderParams: - """Return params effective for the given sigma (getter; single source of truth).""" - sigma_val = float(sigma.item() if isinstance(sigma, torch.Tensor) else sigma) - return _params_for_sigma_from_sorted_dict(sigma_val, self._params_by_sigma) - - def build_from_sigma(self, sigma: float | torch.Tensor) -> MultiModalGuider: - """Return a MultiModalGuider with params effective for the given sigma.""" - return MultiModalGuider( - params=self.params(sigma), - negative_context=self.negative_context, - ) - - -def create_multimodal_guider_factory( - params: MultiModalGuiderParams | MultiModalGuiderFactory, - negative_context: torch.Tensor | None = None, -) -> MultiModalGuiderFactory: - """ - Create or return a MultiModalGuiderFactory. Pass constant params for a - single-params factory (uses MultiModalGuiderFactory.constant), or an existing - MultiModalGuiderFactory. When given a factory, returns it as-is unless - negative_context is provided. For sigma-dependent params use - MultiModalGuiderFactory.from_dict(...) and pass that as params. - """ - if isinstance(params, MultiModalGuiderFactory): - if negative_context is not None and params.negative_context is not negative_context: - return MultiModalGuiderFactory.from_dict(dict(params._params_by_sigma), negative_context=negative_context) - return params - return MultiModalGuiderFactory.constant(params, negative_context=negative_context) - - -def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor: - batch_size = to_project.shape[0] - positive_flat = to_project.reshape(batch_size, -1) - negative_flat = project_onto.reshape(batch_size, -1) - dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) - squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 - return dot_product / squared_norm diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/noisers.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/noisers.py deleted file mode 100644 index fdbebd9ab..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/noisers.py +++ /dev/null @@ -1,36 +0,0 @@ -from dataclasses import replace -from typing import Protocol - -import torch -from ltx_core.types import LatentState - - -class Noiser(Protocol): - """Protocol for adding noise to a latent state during diffusion.""" - - def __call__(self, latent_state: LatentState, noise_scale: float) -> LatentState: ... - - -class GaussianNoiser(Noiser): - """Adds Gaussian noise to a latent state, scaled by the denoise mask.""" - - def __init__(self, generator: torch.Generator): - super().__init__() - self.generator = generator - - def _sample_noise(self, latent_state: LatentState) -> torch.Tensor: - return torch.randn( - *latent_state.latent.shape, - device=latent_state.latent.device, - dtype=latent_state.latent.dtype, - generator=self.generator, - ) - - def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState: - noise = self._sample_noise(latent_state) - latent = torch.lerp(latent_state.latent.float(), noise.float(), noise_scale) - latent = torch.lerp(latent_state.clean_latent.float(), latent, latent_state.denoise_mask) - return replace( - latent_state, - latent=latent.to(latent_state.latent.dtype), - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/patchifiers.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/patchifiers.py deleted file mode 100644 index b46e7bf0c..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/patchifiers.py +++ /dev/null @@ -1,352 +0,0 @@ -import math -from typing import Optional, Tuple - -import einops -import torch -from ltx_core.components.protocols import Patchifier -from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape - - -class VideoLatentPatchifier(Patchifier): - def __init__(self, patch_size: int): - # Patch sizes for video latents. - self._patch_size = ( - 1, # temporal dimension - patch_size, # height dimension - patch_size, # width dimension - ) - - @property - def patch_size(self) -> Tuple[int, int, int]: - return self._patch_size - - def get_token_count(self, tgt_shape: VideoLatentShape) -> int: - return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size) - - def patchify( - self, - latents: torch.Tensor, - ) -> torch.Tensor: - latents = einops.rearrange( - latents, - "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", - p1=self._patch_size[0], - p2=self._patch_size[1], - p3=self._patch_size[2], - ) - - return latents - - def unpatchify( - self, - latents: torch.Tensor, - output_shape: VideoLatentShape, - ) -> torch.Tensor: - assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier" - - patch_grid_frames = output_shape.frames // self._patch_size[0] - patch_grid_height = output_shape.height // self._patch_size[1] - patch_grid_width = output_shape.width // self._patch_size[2] - - latents = einops.rearrange( - latents, - "b (f h w) (c p q) -> b c f (h p) (w q)", - f=patch_grid_frames, - h=patch_grid_height, - w=patch_grid_width, - p=self._patch_size[1], - q=self._patch_size[2], - ) - - return latents - - def get_patch_grid_bounds( - self, - output_shape: AudioLatentShape | VideoLatentShape, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - """ - Return the per-dimension bounds [inclusive start, exclusive end) for every - patch produced by `patchify`. The bounds are expressed in the original - video grid coordinates: frame/time, height, and width. - The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where: - - axis 1 (size 3) enumerates (frame/time, height, width) dimensions - - axis 3 (size 2) stores `[start, end)` indices within each dimension - Args: - output_shape: Video grid description containing frames, height, and width. - device: Device of the latent tensor. - """ - if not isinstance(output_shape, VideoLatentShape): - raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates") - - frames = output_shape.frames - height = output_shape.height - width = output_shape.width - batch_size = output_shape.batch - - # Validate inputs to ensure positive dimensions - assert frames > 0, f"frames must be positive, got {frames}" - assert height > 0, f"height must be positive, got {height}" - assert width > 0, f"width must be positive, got {width}" - assert batch_size > 0, f"batch_size must be positive, got {batch_size}" - - # Generate grid coordinates for each dimension (frame, height, width) - # We use torch.arange to create the starting coordinates for each patch. - # indexing='ij' ensures the dimensions are in the order (frame, height, width). - grid_coords = torch.meshgrid( - torch.arange(start=0, end=frames, step=self._patch_size[0], device=device), - torch.arange(start=0, end=height, step=self._patch_size[1], device=device), - torch.arange(start=0, end=width, step=self._patch_size[2], device=device), - indexing="ij", - ) - - # Stack the grid coordinates to create the start coordinates tensor. - # Shape becomes (3, grid_f, grid_h, grid_w) - patch_starts = torch.stack(grid_coords, dim=0) - - # Create a tensor containing the size of a single patch: - # (frame_patch_size, height_patch_size, width_patch_size). - # Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates. - patch_size_delta = torch.tensor( - self._patch_size, - device=patch_starts.device, - dtype=patch_starts.dtype, - ).view(3, 1, 1, 1) - - # Calculate end coordinates: start + patch_size - # Shape becomes (3, grid_f, grid_h, grid_w) - patch_ends = patch_starts + patch_size_delta - - # Stack start and end coordinates together along the last dimension - # Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end] - latent_coords = torch.stack((patch_starts, patch_ends), dim=-1) - - # Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence. - # Final Shape: (batch_size, 3, num_patches, 2) - latent_coords = einops.repeat( - latent_coords, - "c f h w bounds -> b c (f h w) bounds", - b=batch_size, - bounds=2, - ) - - return latent_coords - - -def get_pixel_coords( - latent_coords: torch.Tensor, - scale_factors: SpatioTemporalScaleFactors, - causal_fix: bool = False, -) -> torch.Tensor: - """ - Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling - each axis (frame/time, height, width) with the corresponding VAE downsampling factors. - Optionally compensate for causal encoding that keeps the first frame at unit temporal scale. - Args: - latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`. - scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied - per axis. - causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs - that treat frame zero differently still yield non-negative timestamps. - """ - # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout. - # Axis 1 of `latent_coords` is ordered (frame/time, height, width) — match that explicitly by - # pulling fields from the NamedTuple rather than relying on tuple iteration order. - broadcast_shape = [1] * latent_coords.ndim - broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width) - scale_tensor = torch.tensor( - [scale_factors.time, scale_factors.height, scale_factors.width], - device=latent_coords.device, - ).view(*broadcast_shape) - - # Apply per-axis scaling to convert latent bounds into pixel-space coordinates. - pixel_coords = latent_coords * scale_tensor - - if causal_fix: - # VAE temporal stride for the very first frame is 1 instead of `scale_factors.time`. - # Shift and clamp to keep the first-frame timestamps causal and non-negative. - pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors.time).clamp(min=0) - - return pixel_coords - - -class AudioPatchifier(Patchifier): - def __init__( - self, - patch_size: int, - sample_rate: int = 16000, - hop_length: int = 160, - audio_latent_downsample_factor: int = 4, - is_causal: bool = True, - shift: int = 0, - ): - """ - Patchifier tailored for spectrogram/audio latents. - Args: - patch_size: Number of mel bins combined into a single patch. This - controls the resolution along the frequency axis. - sample_rate: Original waveform sampling rate. Used to map latent - indices back to seconds so downstream consumers can align audio - and video cues. - hop_length: Window hop length used for the spectrogram. Determines - how many real-time samples separate two consecutive latent frames. - audio_latent_downsample_factor: Ratio between spectrogram frames and - latent frames; compensates for additional downsampling inside the - VAE encoder. - is_causal: When True, timing is shifted to account for causal - receptive fields so timestamps do not peek into the future. - shift: Integer offset applied to the latent indices. Enables - constructing overlapping windows from the same latent sequence. - """ - self.hop_length = hop_length - self.sample_rate = sample_rate - self.audio_latent_downsample_factor = audio_latent_downsample_factor - self.is_causal = is_causal - self.shift = shift - self._patch_size = (1, patch_size, patch_size) - - @property - def patch_size(self) -> Tuple[int, int, int]: - return self._patch_size - - def get_token_count(self, tgt_shape: AudioLatentShape) -> int: - return tgt_shape.frames - - def _get_audio_latent_time_in_sec( - self, - start_latent: int, - end_latent: int, - dtype: torch.dtype, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - """ - Converts latent indices into real-time seconds while honoring causal - offsets and the configured hop length. - Args: - start_latent: Inclusive start index inside the latent sequence. This - sets the first timestamp returned. - end_latent: Exclusive end index. Determines how many timestamps get - generated. - dtype: Floating-point dtype used for the returned tensor, allowing - callers to control precision. - device: Target device for the timestamp tensor. When omitted the - computation occurs on CPU to avoid surprising GPU allocations. - """ - if device is None: - device = torch.device("cpu") - - audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device) - - audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor - - if self.is_causal: - # Frame offset for causal alignment. - # The "+1" ensures the timestamp corresponds to the first sample that is fully available. - causal_offset = 1 - audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0) - - return audio_mel_frame * self.hop_length / self.sample_rate - - def _compute_audio_timings( - self, - batch_size: int, - num_steps: int, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - """ - Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame. - This helper method underpins `get_patch_grid_bounds` for the audio patchifier. - Args: - batch_size: Number of sequences to broadcast the timings over. - num_steps: Number of latent frames (time steps) to convert into timestamps. - device: Device on which the resulting tensor should reside. - """ - resolved_device = device - if resolved_device is None: - resolved_device = torch.device("cpu") - - start_timings = self._get_audio_latent_time_in_sec( - self.shift, - num_steps + self.shift, - torch.float32, - resolved_device, - ) - start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1) - - end_timings = self._get_audio_latent_time_in_sec( - self.shift + 1, - num_steps + self.shift + 1, - torch.float32, - resolved_device, - ) - end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1) - - return torch.stack([start_timings, end_timings], dim=-1) - - def patchify( - self, - audio_latents: torch.Tensor, - ) -> torch.Tensor: - """ - Flattens the audio latent tensor along time. Use `get_patch_grid_bounds` - to derive timestamps for each latent frame based on the configured hop - length and downsampling. - Args: - audio_latents: Latent tensor to patchify. - Returns: - Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the - corresponding timing metadata when needed. - """ - audio_latents = einops.rearrange( - audio_latents, - "b c t f -> b t (c f)", - ) - - return audio_latents - - def unpatchify( - self, - audio_latents: torch.Tensor, - output_shape: AudioLatentShape, - ) -> torch.Tensor: - """ - Restores the `(B, C, T, F)` spectrogram tensor from flattened patches. - Use `get_patch_grid_bounds` to recompute the timestamps that describe each - frame's position in real time. - Args: - audio_latents: Latent tensor to unpatchify. - output_shape: Shape of the unpatched output tensor. - Returns: - Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing - metadata associated with the restored latents. - """ - # audio_latents shape: (batch, time, freq * channels) - audio_latents = einops.rearrange( - audio_latents, - "b t (c f) -> b c t f", - c=output_shape.channels, - f=output_shape.mel_bins, - ) - - return audio_latents - - def get_patch_grid_bounds( - self, - output_shape: AudioLatentShape | VideoLatentShape, - device: Optional[torch.device] = None, - ) -> torch.Tensor: - """ - Return the temporal bounds `[inclusive start, exclusive end)` for every - patch emitted by `patchify`. For audio this corresponds to timestamps in - seconds aligned with the original spectrogram grid. - The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where: - - axis 1 (size 1) represents the temporal dimension - - axis 3 (size 2) stores the `[start, end)` timestamps per patch - Args: - output_shape: Audio grid specification describing the number of time steps. - device: Target device for the returned tensor. - """ - if not isinstance(output_shape, AudioLatentShape): - raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates") - - return self._compute_audio_timings(output_shape.batch, output_shape.frames, device) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/protocols.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/protocols.py deleted file mode 100644 index 40f22a0ff..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/protocols.py +++ /dev/null @@ -1,98 +0,0 @@ -from typing import Protocol, Tuple - -import torch -from ltx_core.types import AudioLatentShape, VideoLatentShape - - -class Patchifier(Protocol): - """ - Protocol for patchifiers that convert latent tensors into patches and assemble them back. - """ - - def patchify( - self, - latents: torch.Tensor, - ) -> torch.Tensor: - ... - """ - Convert latent tensors into flattened patch tokens. - Args: - latents: Latent tensor to patchify. - Returns: - Flattened patch tokens tensor. - """ - - def unpatchify( - self, - latents: torch.Tensor, - output_shape: AudioLatentShape | VideoLatentShape, - ) -> torch.Tensor: - """ - Converts latent tensors between spatio-temporal formats and flattened sequence representations. - Args: - latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`. - output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or - VideoLatentShape. - Returns: - Dense latent tensor restored from the flattened representation. - """ - - @property - def patch_size(self) -> Tuple[int, int, int]: - ... - """ - Returns the patch size as a tuple of (temporal, height, width) dimensions - """ - - def get_patch_grid_bounds( - self, - output_shape: AudioLatentShape | VideoLatentShape, - device: torch.device | None = None, - ) -> torch.Tensor: - ... - """ - Compute metadata describing where each latent patch resides within the - grid specified by `output_shape`. - Args: - output_shape: Target grid layout for the patches. - device: Target device for the returned tensor. - Returns: - Tensor containing patch coordinate metadata such as spatial or temporal intervals. - """ - - -class SchedulerProtocol(Protocol): - """ - Protocol for schedulers that provide a sigmas schedule tensor for a - given number of steps. Device is cpu. - """ - - def execute(self, steps: int, **kwargs) -> torch.FloatTensor: ... - - -class GuiderProtocol(Protocol): - """ - Protocol for guiders that compute a delta tensor given conditioning inputs. - The returned delta should be added to the conditional output (cond), enabling - multiple guiders to be chained together by accumulating their deltas. - """ - - scale: float - - def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: ... - - def enabled(self) -> bool: - """ - Returns whether the corresponding perturbation is enabled. E.g. for CFG, this should return False if the scale - is 1.0. - """ - ... - - -class DiffusionStepProtocol(Protocol): - """ - Protocol for diffusion steps that provide a next sample tensor for a given current sample tensor, - current denoised sample tensor, and sigmas tensor. - """ - - def step(self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **kwargs) -> torch.Tensor: ... diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/schedulers.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/schedulers.py deleted file mode 100644 index 3e395f4c4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/schedulers.py +++ /dev/null @@ -1,125 +0,0 @@ -import math -from functools import lru_cache - -import numpy -import scipy -import torch -from ltx_core.components.protocols import SchedulerProtocol - -BASE_SHIFT_ANCHOR = 1024 -MAX_SHIFT_ANCHOR = 4096 - - -class LTX2Scheduler(SchedulerProtocol): - """ - Default scheduler for LTX-2 diffusion sampling. - Generates a sigma schedule with token-count-dependent shifting and optional - stretching to a terminal value. - """ - - def execute( - self, - steps: int, - latent: torch.Tensor | None = None, - max_shift: float = 2.05, - base_shift: float = 0.95, - stretch: bool = True, - terminal: float = 0.1, - default_number_of_tokens: int = MAX_SHIFT_ANCHOR, - **_kwargs, - ) -> torch.FloatTensor: - tokens = math.prod(latent.shape[2:]) if latent is not None else default_number_of_tokens - sigmas = torch.linspace(1.0, 0.0, steps + 1) - - x1 = BASE_SHIFT_ANCHOR - x2 = MAX_SHIFT_ANCHOR - mm = (max_shift - base_shift) / (x2 - x1) - b = base_shift - mm * x1 - sigma_shift = (tokens) * mm + b - - power = 1 - sigmas = torch.where( - sigmas != 0, - math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power), - 0, - ) - - # Stretch sigmas so that its final value matches the given terminal value. - if stretch: - non_zero_mask = sigmas != 0 - non_zero_sigmas = sigmas[non_zero_mask] - one_minus_z = 1.0 - non_zero_sigmas - scale_factor = one_minus_z[-1] / (1.0 - terminal) - stretched = 1.0 - (one_minus_z / scale_factor) - sigmas[non_zero_mask] = stretched - - return sigmas.to(torch.float32) - - -class LinearQuadraticScheduler(SchedulerProtocol): - """ - Scheduler with linear steps followed by quadratic steps. - Produces a sigma schedule that transitions linearly up to a threshold, - then follows a quadratic curve for the remaining steps. - """ - - def execute(self, steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None, **_kwargs) -> torch.FloatTensor: - if steps == 1: - return torch.FloatTensor([1.0, 0.0]) - - if linear_steps is None: - linear_steps = steps // 2 - linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] - threshold_noise_step_diff = linear_steps - threshold_noise * steps - quadratic_steps = steps - linear_steps - quadratic_sigma_schedule = [] - if quadratic_steps > 0: - quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) - linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) - const = quadratic_coef * (linear_steps**2) - quadratic_sigma_schedule = [quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)] - sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0] - sigma_schedule = [1.0 - x for x in sigma_schedule] - return torch.FloatTensor(sigma_schedule) - - -class BetaScheduler(SchedulerProtocol): - """ - Scheduler using a beta distribution to sample timesteps. - Based on: https://arxiv.org/abs/2407.12173 - """ - - shift = 2.37 - timesteps_length = 10000 - - def execute(self, steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.FloatTensor: - """ - Execute the beta scheduler. - Args: - steps: The number of steps to execute the scheduler for. - alpha: The alpha parameter for the beta distribution. - beta: The beta parameter for the beta distribution. - Warnings: - The number of steps within `sigmas` theoretically might be less than `steps+1`, - because of the deduplication of the identical timesteps - Returns: - A tensor of sigmas. - """ - model_sampling_sigmas = _precalculate_model_sampling_sigmas(self.shift, self.timesteps_length) - total_timesteps = len(model_sampling_sigmas) - 1 - ts = 1 - numpy.linspace(0, 1, steps, endpoint=False) - ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps).tolist() - ts = list(dict.fromkeys(ts)) - - sigmas = [float(model_sampling_sigmas[int(t)]) for t in ts] + [0.0] - return torch.FloatTensor(sigmas) - - -@lru_cache(maxsize=5) -def _precalculate_model_sampling_sigmas(shift: float, timesteps_length: int) -> torch.Tensor: - timesteps = torch.arange(1, timesteps_length + 1, 1) / timesteps_length - return torch.Tensor([flux_time_shift(shift, 1.0, t) for t in timesteps]) - - -def flux_time_shift(mu: float, sigma: float, t: float) -> float: - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/__init__.py deleted file mode 100644 index 8d3eb9f2e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Conditioning utilities: latent state, tools, and conditioning types.""" - -from ltx_core.conditioning.exceptions import ConditioningError -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.conditioning.types import ( - AudioConditionByReferenceLatent, - ConditioningItemAttentionStrengthWrapper, - VideoConditionByKeyframeIndex, - VideoConditionByLatentIndex, - VideoConditionByMask, - VideoConditionByReferenceLatent, -) - -__all__ = [ - "AudioConditionByReferenceLatent", - "ConditioningError", - "ConditioningItem", - "ConditioningItemAttentionStrengthWrapper", - "VideoConditionByKeyframeIndex", - "VideoConditionByLatentIndex", - "VideoConditionByMask", - "VideoConditionByReferenceLatent", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/exceptions.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/exceptions.py deleted file mode 100644 index 458aa3b47..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/exceptions.py +++ /dev/null @@ -1,4 +0,0 @@ -class ConditioningError(Exception): - """ - Class for conditioning-related errors. - """ diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/item.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/item.py deleted file mode 100644 index e6eef52d1..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/item.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import Protocol - -from ltx_core.tools import LatentTools -from ltx_core.types import LatentState - - -class ConditioningItem(Protocol): - """Protocol for conditioning items that modify latent state during diffusion.""" - - def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: - """ - Apply the conditioning to the latent state. - Args: - latent_state: The latent state to apply the conditioning to. This is state always patchified. - Returns: - The latent state after the conditioning has been applied. - IMPORTANT: If the conditioning needs to add extra tokens to the latent, it should add them to the end of the - latent. - """ - ... diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/mask_utils.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/mask_utils.py deleted file mode 100644 index 4833cc6df..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/mask_utils.py +++ /dev/null @@ -1,200 +0,0 @@ -"""Utilities for building 2D self-attention masks for conditioning items.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from ltx_core.types import LatentState - - -def resolve_cross_mask( - attention_mask: float | int | torch.Tensor, - num_new_tokens: int, - batch_size: int, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - """Convert an attention_mask (scalar or tensor) to a (B, M) cross_mask tensor. - Args: - attention_mask: Scalar value applied uniformly, 1D tensor of shape (M,) - broadcast across batch, or 2D tensor of shape (B, M). - num_new_tokens: Number of new conditioning tokens M. - batch_size: Batch size B. - device: Device for the output tensor. - dtype: Data type for the output tensor. - Returns: - Cross-mask tensor of shape (B, M). - """ - if isinstance(attention_mask, (int, float)): - return torch.full( - (batch_size, num_new_tokens), - fill_value=float(attention_mask), - device=device, - dtype=dtype, - ) - mask = attention_mask.to(device=device, dtype=dtype) - - # Handle scalar (0-D) tensor like a Python scalar. - if mask.dim() == 0: - return torch.full( - (batch_size, num_new_tokens), - fill_value=float(mask.item()), - device=device, - dtype=dtype, - ) - - if mask.dim() == 1: - if mask.shape[0] != num_new_tokens: - raise ValueError(f"1-D attention_mask length must equal num_new_tokens ({num_new_tokens}), got shape {tuple(mask.shape)}") - mask = mask.unsqueeze(0).expand(batch_size, -1) - elif mask.dim() == 2: - b, m = mask.shape - if m != num_new_tokens: - raise ValueError(f"2-D attention_mask second dimension must equal num_new_tokens ({num_new_tokens}), got shape {tuple(mask.shape)}") - if b not in (batch_size, 1): - raise ValueError(f"2-D attention_mask batch dimension must equal batch_size ({batch_size}) or 1, got shape {tuple(mask.shape)}") - if b == 1 and batch_size > 1: - mask = mask.expand(batch_size, -1) - else: - raise ValueError(f"attention_mask tensor must be 0-D, 1-D, or 2-D, got {mask.dim()}-D with shape {tuple(mask.shape)}") - return mask - - -def update_attention_mask( - latent_state: LatentState, - attention_mask: float | torch.Tensor | None, - num_noisy_tokens: int, - num_new_tokens: int, - batch_size: int, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor | None: - """Build or update the self-attention mask for newly appended conditioning tokens. - If *attention_mask* is ``None`` and no existing mask is present, returns - ``None``. If *attention_mask* is ``None`` but an existing mask is present, - the mask is expanded with full attention (1s) for the new tokens so that - its dimensions stay consistent with the growing latent sequence. Otherwise, - resolves *attention_mask* to a per-token cross-mask and expands the 2-D - attention mask via :func:`build_attention_mask`. - Args: - latent_state: Current latent state (provides the existing mask and total - existing-token count). - attention_mask: Per-token attention weight. Scalar, 1-D ``(M,)``, 2-D - ``(B, M)`` tensor, or ``None`` (no-op). - num_noisy_tokens: Number of original noisy tokens (from - ``latent_tools.target_shape.token_count()``). - num_new_tokens: Number of new conditioning tokens being appended. - batch_size: Batch size. - device: Device for the output tensor. - dtype: Data type for the output tensor. - Returns: - Updated attention mask of shape ``(B, N+M, N+M)``, or ``None`` if no - masking is needed. - """ - if attention_mask is None: - if latent_state.attention_mask is None: - return None - # Existing mask present but no new mask requested: pad with 1s (full - # attention) so the mask dimensions stay consistent with the growing - # latent sequence. - cross_mask = torch.ones(batch_size, num_new_tokens, device=device, dtype=dtype) - return build_attention_mask( - existing_mask=latent_state.attention_mask, - num_noisy_tokens=num_noisy_tokens, - num_new_tokens=num_new_tokens, - num_existing_tokens=latent_state.latent.shape[1], - cross_mask=cross_mask, - device=device, - dtype=dtype, - ) - - cross_mask = resolve_cross_mask(attention_mask, num_new_tokens, batch_size, device, dtype) - return build_attention_mask( - existing_mask=latent_state.attention_mask, - num_noisy_tokens=num_noisy_tokens, - num_new_tokens=num_new_tokens, - num_existing_tokens=latent_state.latent.shape[1], - cross_mask=cross_mask, - device=device, - dtype=dtype, - ) - - -def build_attention_mask( - existing_mask: torch.Tensor | None, - num_noisy_tokens: int, - num_new_tokens: int, - num_existing_tokens: int, - cross_mask: torch.Tensor, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - """ - Expand the attention mask to include newly appended conditioning tokens. - Each conditioning item appends M new reference tokens to the sequence. This function - builds a (B, N+M, N+M) attention mask with the following block structure: - noisy prev_ref new_ref - (N_noisy) (N-N_noisy) (M) - ┌───────────┬───────────┬───────────┐ - noisy │ │ │ │ - (N_noisy) │ existing │ existing │ cross │ - │ │ │ │ - ├───────────┼───────────┼───────────┤ - prev_ref │ │ │ │ - (N-N_noisy)│ existing │ existing │ 0 │ - │ │ │ │ - ├───────────┼───────────┼───────────┤ - new_ref │ │ │ │ - (M) │ cross │ 0 │ 1 │ - │ │ │ │ - └───────────┴───────────┴───────────┘ - Where: - - **existing**: preserved from the previous mask (or 1.0 if first conditioning) - - **cross**: values from *cross_mask* (shape B, M), in [0, 1] - - **0**: no attention between different reference groups - Args: - existing_mask: Current attention mask of shape (B, N, N), or None if no mask exists yet. - When None, the top-left NxN block is filled with 1s (full attention between all - existing tokens including any prior reference tokens that had no mask). - num_noisy_tokens: Number of original noisy tokens (always at positions [0:num_noisy_tokens]). - num_new_tokens: Number of new conditioning tokens M being appended. - num_existing_tokens: Total number of current tokens N (noisy + any prior conditioning tokens). - cross_mask: Per-token attention weight of shape (B, M) controlling attention between - new reference tokens and noisy tokens. Values in [0, 1]. - device: Device for the output tensor. - dtype: Data type for the output tensor. - Returns: - Attention mask of shape (B, N+M, N+M) with values in [0, 1]. - """ - batch_size = cross_mask.shape[0] - total = num_existing_tokens + num_new_tokens - - # Start with zeros - mask = torch.zeros((batch_size, total, total), device=device, dtype=dtype) - - # Top-left: preserve existing mask or fill with 1s for noisy tokens - if existing_mask is not None: - mask[:, :num_existing_tokens, :num_existing_tokens] = existing_mask - else: - mask[:, :num_existing_tokens, :num_existing_tokens] = 1.0 - - # Bottom-right: new reference tokens fully attend to themselves - mask[:, num_existing_tokens:, num_existing_tokens:] = 1.0 - - # Cross-attention between noisy tokens and new reference tokens - # cross_mask shape: (B, M) -> broadcast to (B, N_noisy, M) and (B, M, N_noisy) - - # Noisy tokens attending to new reference tokens: [0:N_noisy, N:N+M] - # Each column j in this block gets cross_mask[:, j] - mask[:, :num_noisy_tokens, num_existing_tokens:] = cross_mask.unsqueeze(1) - - # New reference tokens attending to noisy tokens: [N:N+M, 0:N_noisy] - # Each row i in this block gets cross_mask[:, i] - mask[:, num_existing_tokens:, :num_noisy_tokens] = cross_mask.unsqueeze(2) - - # [N_noisy:N, N:N+M] and [N:N+M, N_noisy:N] remain 0 (no cross-ref attention) - - return mask diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/__init__.py deleted file mode 100644 index bd8d96227..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Conditioning type implementations.""" - -from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper -from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex -from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex -from ltx_core.conditioning.types.mask_cond import VideoConditionByMask -from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent -from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent - -__all__ = [ - "AudioConditionByReferenceLatent", - "ConditioningItemAttentionStrengthWrapper", - "VideoConditionByKeyframeIndex", - "VideoConditionByLatentIndex", - "VideoConditionByMask", - "VideoConditionByReferenceLatent", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/attention_strength_wrapper.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/attention_strength_wrapper.py deleted file mode 100644 index a5d1f495b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/attention_strength_wrapper.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Wrapper conditioning item that adds attention masking to any inner conditioning.""" - -from dataclasses import replace - -import torch -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.conditioning.mask_utils import update_attention_mask -from ltx_core.tools import LatentTools -from ltx_core.types import LatentState - - -class ConditioningItemAttentionStrengthWrapper(ConditioningItem): - """Wraps a conditioning item to add an attention mask for its tokens. - Separates the *attention-masking* concern from the underlying conditioning - logic (token layout, positional encoding, denoise strength). The inner - conditioning item appends tokens to the latent sequence as usual, and this - wrapper then builds or updates the self-attention mask so that the newly - added tokens interact with the noisy tokens according to *attention_mask*. - Args: - conditioning: Any conditioning item that appends tokens to the latent. - attention_mask: Per-token attention weight controlling how strongly the - new conditioning tokens attend to/from noisy tokens. Can be a - scalar (float) applied uniformly, or a tensor of shape ``(B, M)`` - for spatial control, where ``M = F * H * W`` is the number of - patchified conditioning tokens. Values in ``[0, 1]``. - Example:: - cond = ConditioningItemAttentionStrengthWrapper( - VideoConditionByReferenceLatent(latent=ref, strength=1.0), - attention_mask=0.5, - ) - state = cond.apply_to(latent_state, latent_tools) - """ - - def __init__( - self, - conditioning: ConditioningItem, - attention_mask: float | torch.Tensor, - ): - self.conditioning = conditioning - self.attention_mask = attention_mask - - def apply_to( - self, - latent_state: LatentState, - latent_tools: LatentTools, - ) -> LatentState: - """Apply inner conditioning, then build the attention mask for its tokens.""" - # Snapshot the original state for mask building - original_state = latent_state - - # Inner conditioning appends tokens (positions, denoise mask, etc.) - new_state = self.conditioning.apply_to(latent_state, latent_tools) - - num_new_tokens = new_state.latent.shape[1] - original_state.latent.shape[1] - if num_new_tokens == 0: - return new_state - - # Build the attention mask using the *original* state as the reference - # so that the block structure is computed correctly. - new_attention_mask = update_attention_mask( - latent_state=original_state, - attention_mask=self.attention_mask, - num_noisy_tokens=latent_tools.target_shape.token_count(), - num_new_tokens=num_new_tokens, - batch_size=new_state.latent.shape[0], - device=new_state.latent.device, - dtype=new_state.latent.dtype, - ) - - return replace(new_state, attention_mask=new_attention_mask) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/keyframe_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/keyframe_cond.py deleted file mode 100644 index 03baea429..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/keyframe_cond.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -from ltx_core.components.patchifiers import get_pixel_coords -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.conditioning.mask_utils import update_attention_mask -from ltx_core.tools import VideoLatentTools -from ltx_core.types import LatentState, VideoLatentShape - - -class VideoConditionByKeyframeIndex(ConditioningItem): - """ - Conditions video generation on keyframe latents at a specific frame index. - Appends keyframe tokens to the sequence with positions offset by frame_idx: the keyframe - latents become clean-latent tokens (placeholder zeros in the noisy latent) and the denoise - mask is set from the strength parameter. - To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`. - Args: - keyframes: Keyframe latents [B, C, F, H, W]. - frame_idx: Frame index offset for positional encoding. - strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised). - num_pixel_frames: Number of pixel frames the keyframe latent originally encodes. - """ - - def __init__( - self, - keyframes: torch.Tensor, - frame_idx: int, - strength: float, - num_pixel_frames: int = 1, - ): - self.keyframes = keyframes - self.frame_idx = frame_idx - self.strength = strength - self.num_pixel_frames = num_pixel_frames - - def apply_to( - self, - latent_state: LatentState, - latent_tools: VideoLatentTools, - ) -> LatentState: - tokens = latent_tools.patchifier.patchify(self.keyframes) - latent_coords = latent_tools.patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape.from_torch_shape(self.keyframes.shape), - device=self.keyframes.device, - ) - positions = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=latent_tools.scale_factors, - causal_fix=latent_tools.causal_fix if self.frame_idx == 0 else False, - ) - - positions[:, 0, ...] += self.frame_idx - # If the keyframe latent encodes a single pixel frame, - # narrow the temporal end to [start, start + 1) instead of the - # VAE-scaled range. - if self.num_pixel_frames == 1: - positions[:, 0, ..., 1:] = positions[:, 0, ..., :1] + 1 - positions = positions.to(dtype=torch.float32) - positions[:, 0, ...] /= latent_tools.fps - - denoise_mask = torch.full( - size=(*tokens.shape[:2], 1), - fill_value=1.0 - self.strength, - device=self.keyframes.device, - dtype=self.keyframes.dtype, - ) - - new_attention_mask = update_attention_mask( - latent_state=latent_state, - attention_mask=None, - num_noisy_tokens=latent_tools.target_shape.token_count(), - num_new_tokens=tokens.shape[1], - batch_size=tokens.shape[0], - device=self.keyframes.device, - dtype=self.keyframes.dtype, - ) - - return LatentState( - latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1), - denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1), - positions=torch.cat([latent_state.positions, positions], dim=2), - clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1), - attention_mask=new_attention_mask, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/latent_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/latent_cond.py deleted file mode 100644 index 96f556f36..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/latent_cond.py +++ /dev/null @@ -1,40 +0,0 @@ -import torch -from ltx_core.conditioning.exceptions import ConditioningError -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.tools import LatentTools -from ltx_core.types import LatentState - - -class VideoConditionByLatentIndex(ConditioningItem): - """ - Conditions video generation by injecting latents at a specific latent frame index. - Sets the clean latents at positions corresponding to latent_idx to the injected latents, - sets denoise strength according to the strength parameter. - """ - - def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int): - self.latent = latent - self.strength = strength - self.latent_idx = latent_idx - - def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: - cond_batch, cond_channels, _, cond_height, cond_width = self.latent.shape - tgt_batch, tgt_channels, tgt_frames, tgt_height, tgt_width = latent_tools.target_shape.to_torch_shape() - - if (cond_batch, cond_channels, cond_height, cond_width) != (tgt_batch, tgt_channels, tgt_height, tgt_width): - raise ConditioningError( - f"Can't apply image conditioning item to latent with shape {latent_tools.target_shape}, expected " - f"shape is ({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_height}, {tgt_width}). Make sure " - "the image and latent have the same spatial shape." - ) - - tokens = latent_tools.patchifier.patchify(self.latent) - start_token = latent_tools.patchifier.get_token_count(latent_tools.target_shape._replace(frames=self.latent_idx)) - stop_token = start_token + tokens.shape[1] - - latent_state = latent_state.clone() - - latent_state.clean_latent[:, start_token:stop_token] = tokens - latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength - - return latent_state diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/mask_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/mask_cond.py deleted file mode 100644 index b9fc5541d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/mask_cond.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Mask-based conditioning for inpainting and spatial conditioning.""" - -from dataclasses import replace - -import torch -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.tools import LatentTools -from ltx_core.types import LatentState - - -class VideoConditionByMask(ConditioningItem): - """Condition video generation using a binary mask over latent frames. - Masked positions (mask=1) receive the provided clean latent values and are - excluded from denoising (denoise_mask set to ``1 - strength``). Unmasked - positions (mask=0) are left unchanged and denoised normally. - The mask operates in **unpatchified latent** space — it should have shape - ``[B, F, H, W]`` matching the latent dimensions (after VAE encoding, - before patchification). This is consistent with the latent input format - used by all other conditioning items. - Args: - latent: Clean conditioning latents in unpatchified format [B, C, F, H, W]. - Must match the target shape of the latent tools. - mask: Binary mask [B, F, H, W] in unpatchified latent space. - 1 = conditioning position (clean, excluded from denoising), - 0 = generated position (noised, denoised normally). - strength: Conditioning strength for masked positions. 1.0 = fully clean - (no denoising), 0.0 = no conditioning effect. Default 1.0. - """ - - def __init__(self, latent: torch.Tensor, mask: torch.Tensor, strength: float = 1.0): - self.latent = latent - self.mask = mask - self.strength = strength - - def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: - """Apply mask-based conditioning to the latent state.""" - tokens = latent_tools.patchifier.patchify(self.latent) - - mask = latent_tools.patchifier.patchify(self.mask.unsqueeze(1)) - - m = mask.to(dtype=latent_state.latent.dtype) - inv = 1 - m - - return replace( - latent_state, - clean_latent=latent_state.clean_latent * inv + tokens * m, - denoise_mask=latent_state.denoise_mask * inv + (1.0 - self.strength) * m, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/noise_mask_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/noise_mask_cond.py deleted file mode 100644 index 62a78d44f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/noise_mask_cond.py +++ /dev/null @@ -1,43 +0,0 @@ -from dataclasses import dataclass - -from ltx_core.components.patchifiers import get_pixel_coords -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.tools import LatentTools, SpatioTemporalScaleFactors -from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape - - -@dataclass(frozen=True) -class TemporalRegionMask(ConditioningItem): - """Conditioning item that sets ``denoise_mask = 0`` outside a time range - and ``1`` inside, so only the specified temporal region is regenerated. - Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified* - (token) space using the patchifier's ``get_patch_grid_bounds``: for video - coords are latent frame indices (converted from seconds via ``fps``), for - audio coords are already in seconds. - """ - - start_time: float # seconds, inclusive - end_time: float # seconds, exclusive - fps: float - - def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: - coords = latent_tools.patchifier.get_patch_grid_bounds(latent_tools.target_shape, device=latent_state.denoise_mask.device) - if isinstance(latent_tools.target_shape, AudioLatentShape): - # Audio: patchifier get_patch_grid_bounds returns seconds - t_boundaries = coords[:, 0] - elif isinstance(latent_tools.target_shape, VideoLatentShape): - # Video: patchifier get_patch_grid_bounds returns latent bounds, converting to frame numbers & pixel bounds - scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default()) - pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True)) - # converting frame numbers to seconds - t_boundaries = pixel_bounds[:, 0] / self.fps - else: - raise ValueError("Unsupported LatentShape type, expected AudioLatentShape or VideoLatentShape") - t_start, t_end = t_boundaries.unbind(dim=-1) # [B, N] - in_region = (t_end > self.start_time) & (t_start < self.end_time) - state = latent_state.clone() - mask_val = in_region.to(state.denoise_mask.dtype) - if state.denoise_mask.dim() == 3: - mask_val = mask_val.unsqueeze(-1) - state.denoise_mask.copy_(mask_val) - return state diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_audio_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_audio_cond.py deleted file mode 100644 index 682d9b7af..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_audio_cond.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Audio reference conditioning items.""" - -from __future__ import annotations - -import torch -from ltx_core.conditioning.mask_utils import update_attention_mask -from ltx_core.tools import LatentTools -from ltx_core.types import LatentState - - -class AudioConditionByReferenceLatent: - """Append patchified reference audio tokens after the target audio sequence. - Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent` - but for audio. The reference tokens are appended so the target audio tokens stay - in the first ``num_noisy_tokens`` positions and can be kept by - :meth:`ltx_core.tools.LatentTools.clear_conditioning`. - Args: - patchified: Patchified reference latent ``[B, T_ref, C]``. - positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``. - strength: 1.0 keeps reference clean; 0.0 would fully denoise it. - """ - - def __init__( - self, - patchified: torch.Tensor, - positions: torch.Tensor, - strength: float = 1.0, - ) -> None: - self.patchified = patchified - self.positions = positions.to(dtype=torch.float32) - self.strength = strength - - def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: - tokens = self.patchified - denoise_mask = torch.full( - size=(*tokens.shape[:2], 1), - fill_value=1.0 - self.strength, - device=tokens.device, - dtype=tokens.dtype, - ) - - new_attention_mask = update_attention_mask( - latent_state=latent_state, - attention_mask=None, - num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape), - num_new_tokens=tokens.shape[1], - batch_size=tokens.shape[0], - device=tokens.device, - dtype=tokens.dtype, - ) - - return LatentState( - latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1), - denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1), - positions=torch.cat([latent_state.positions, self.positions], dim=2), - clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1), - attention_mask=new_attention_mask, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_video_cond.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_video_cond.py deleted file mode 100644 index be0f8f67b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_video_cond.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Reference video conditioning for IC-LoRA inference.""" - -import torch -from ltx_core.components.patchifiers import get_pixel_coords -from ltx_core.conditioning.item import ConditioningItem -from ltx_core.conditioning.mask_utils import update_attention_mask -from ltx_core.tools import VideoLatentTools -from ltx_core.types import LatentState, VideoLatentShape - - -class VideoConditionByReferenceLatent(ConditioningItem): - """ - Conditions video generation on a reference video latent for IC-LoRA inference. - IC-LoRAs are trained by concatenating reference (control signal) and target tokens, - learning to attend across both. This class replicates that setup at inference by - appending the reference tokens to the sequence as clean latents (with placeholder zeros - in the noisy latent). - IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px - reference for 768px output) for efficiency and better generalization. The - `downscale_factor` scales reference positions to match target coordinates, preserving - the learned positional relationships. This must match the factor used during training - (stored in LoRA metadata). - To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`. - Args: - latent: Reference video latents [B, C, F, H, W]. - downscale_factor: Target/reference spatial ratio (e.g. 2 = half-res ref). - temporal_scale_factor: Target/reference temporal ratio S (e.g. 4 = ref at 1/4 fps). - strength: Conditioning strength. 1.0 = full (reference kept clean), - 0.0 = none (reference denoised). Default 1.0. - """ - - def __init__( - self, - latent: torch.Tensor, - downscale_factor: int = 1, - temporal_scale_factor: int = 1, - strength: float = 1.0, - ): - self.latent = latent - self.downscale_factor = downscale_factor - self.temporal_scale_factor = temporal_scale_factor - self.strength = strength - - def apply_to( - self, - latent_state: LatentState, - latent_tools: VideoLatentTools, - ) -> LatentState: - """Append reference video tokens with positions translated into the target frame.""" - tokens = latent_tools.patchifier.patchify(self.latent) - - latent_coords = latent_tools.patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape.from_torch_shape(self.latent.shape), - device=self.latent.device, - ) - positions = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=latent_tools.scale_factors, - causal_fix=latent_tools.causal_fix, - ) - positions = positions.to(dtype=torch.float32) - - # Place ref tokens on their own time spacing (= target_fps / S). - positions[:, 0, ...] /= latent_tools.fps / self.temporal_scale_factor - - # Translate into the target's frame so ref's last patch ends with target's last - # patch; clamp the causal patch's negative start back to [0, 1/target_fps). - if self.temporal_scale_factor != 1: - t_target = latent_state.positions[:, 0, 0:1, 1:2].to(dtype=torch.float32) # = 1/target_fps - positions[:, 0, ...] = torch.clamp( - positions[:, 0, ...] - (self.temporal_scale_factor - 1) * t_target, - min=0, - ) - if self.downscale_factor != 1: - positions[:, 1, ...] *= self.downscale_factor # height axis - positions[:, 2, ...] *= self.downscale_factor # width axis - - denoise_mask = torch.full( - size=(*tokens.shape[:2], 1), - fill_value=1.0 - self.strength, - device=self.latent.device, - dtype=self.latent.dtype, - ) - - new_attention_mask = update_attention_mask( - latent_state=latent_state, - attention_mask=None, - num_noisy_tokens=latent_tools.target_shape.token_count(), - num_new_tokens=tokens.shape[1], - batch_size=tokens.shape[0], - device=self.latent.device, - dtype=self.latent.dtype, - ) - - return LatentState( - latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1), - denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1), - positions=torch.cat([latent_state.positions, positions], dim=2), - clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1), - attention_mask=new_attention_mask, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/__init__.py deleted file mode 100644 index 8214c5c70..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Guidance and perturbation utilities for attention manipulation.""" - -from ltx_core.guidance.perturbations import ( - BatchedPerturbationConfig, - Perturbation, - PerturbationConfig, - PerturbationType, -) - -__all__ = [ - "BatchedPerturbationConfig", - "Perturbation", - "PerturbationConfig", - "PerturbationType", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/perturbations.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/perturbations.py deleted file mode 100644 index 9106a7e8b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/perturbations.py +++ /dev/null @@ -1,77 +0,0 @@ -from dataclasses import dataclass -from enum import Enum - -import torch -from torch._prims_common import DeviceLikeType - - -class PerturbationType(Enum): - """Types of attention perturbations for STG (Spatio-Temporal Guidance).""" - - SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn" - SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn" - SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn" - SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn" - - -@dataclass(frozen=True) -class Perturbation: - """A single perturbation specifying which attention type to skip and in which blocks.""" - - type: PerturbationType - blocks: list[int] | None # None means all blocks - - def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool: - if self.type != perturbation_type: - return False - - if self.blocks is None: - return True - - return block in self.blocks - - -@dataclass(frozen=True) -class PerturbationConfig: - """Configuration holding a list of perturbations for a single sample.""" - - perturbations: list[Perturbation] | None - - def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool: - if self.perturbations is None: - return False - - return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) - - @staticmethod - def empty() -> "PerturbationConfig": - return PerturbationConfig([]) - - -@dataclass(frozen=True) -class BatchedPerturbationConfig: - """Perturbation configurations for a batch, with utilities for generating attention masks.""" - - perturbations: list[PerturbationConfig] - - def mask(self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype) -> torch.Tensor: - mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype) - for batch_idx, perturbation in enumerate(self.perturbations): - if perturbation.is_perturbed(perturbation_type, block): - mask[batch_idx] = 0 - - return mask - - def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor: - mask = self.mask(perturbation_type, block, values.device, values.dtype) - return mask.view(mask.numel(), *([1] * len(values.shape[1:]))) - - def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool: - return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) - - def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool: - return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations) - - @staticmethod - def empty(batch_size: int) -> "BatchedPerturbationConfig": - return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)]) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/hdr.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/hdr.py deleted file mode 100644 index 8413004af..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/hdr.py +++ /dev/null @@ -1,71 +0,0 @@ -"""HDR utilities: LogC3 compression for HDR IC-LoRA training and inference. -Provides compress/decompress and postprocess helpers for HDR video generation. -Used by ltx-pipelines for HDR IC-LoRA and by ltx-trainer for HDR validation. -""" - -from __future__ import annotations - -from typing import Literal - -import torch -from torch import Tensor - - -class LogC3: - """ARRI LogC3 (EI 800) HDR compression. - Maps linear [0, ∞) <-> LogC3 [0, 1] via the camera log curve. The log - curve allocates more precision to shadows/midtones and compresses - highlights smoothly. Callers are responsible for mapping the [0, 1] - output to the VAE's [-1, 1] input range. - """ - - name = "LogC3" - A = 5.555556 - B = 0.052272 - C = 0.247190 - D = 0.385537 - E = 5.367655 - F = 0.092809 - CUT = 0.010591 - - def compress(self, hdr: Tensor) -> Tensor: - """Compress linear HDR [0, ∞) → LogC3 [0, 1].""" - x = torch.clamp(hdr, min=0.0) - log_part = self.C * torch.log10(self.A * x + self.B) + self.D - lin_part = self.E * x + self.F - logc = torch.where(x >= self.CUT, log_part, lin_part) - return torch.clamp(logc, 0.0, 1.0) - - def compress_ldr(self, ldr: Tensor) -> Tensor: - """Compress LDR [0, 1] → [0, 1] (no log curve, just clamp).""" - return torch.clamp(ldr, 0.0, 1.0) - - def decompress(self, logc: Tensor) -> Tensor: - """Decompress LogC3 [0, 1] → linear HDR [0, ∞).""" - logc = torch.clamp(logc, 0.0, 1.0) - cut_log = self.E * self.CUT + self.F - lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A - lin_from_lin = (logc - self.F) / self.E - return torch.where(logc >= cut_log, lin_from_log, lin_from_lin) - - def decompress_ldr(self, logc: Tensor) -> Tensor: - """Decompress [0, 1] → LDR [0, 1] (identity clamp).""" - return torch.clamp(logc, 0.0, 1.0) - - -def apply_hdr_decode_postprocess( - decoded_video: Tensor, - transform: Literal["logc3"] = "logc3", -) -> Tensor: - """Apply HDR decompress to VAE decode output for HDR recovery. - Args: - decoded_video: Tensor from VAE decode in [0, 1], shape [B, C, F, H, W]. - Must be float32 for sufficient color resolution. - transform: "logc3". - Returns: - HDR video tensor float32. - """ - decoded_video = decoded_video.float() - if transform == "logc3": - return LogC3().decompress(decoded_video) - raise ValueError(f"Unsupported HDR transform: {transform}") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/__init__.py deleted file mode 100644 index 9bb099cf5..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/__init__.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Loader utilities for model weights, LoRAs, and safetensor operations.""" - -from ltx_core.loader.fuse_loras import apply_loras -from ltx_core.loader.helpers import ( - create_meta_model, - load_state_dict, - read_model_config, -) -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import ( - LoRAAdaptableProtocol, - LoraPathStrengthAndSDOps, - LoraStateDictWithStrength, - ModelBuilderProtocol, - StateDict, - StateDictLoader, -) -from ltx_core.loader.registry import DummyRegistry, Registry, StateDictRegistry -from ltx_core.loader.sd_ops import ( - LTXV_LORA_COMFY_RENAMING_MAP, - ContentMatching, - ContentReplacement, - KeyValueOperation, - KeyValueOperationResult, - SDKeyValueOperation, - SDOps, -) -from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader, SafetensorsStateDictLoader -from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder - -__all__ = [ - "LTXV_LORA_COMFY_RENAMING_MAP", - "ContentMatching", - "ContentReplacement", - "DummyRegistry", - "KeyValueOperation", - "KeyValueOperationResult", - "LoRAAdaptableProtocol", - "LoraPathStrengthAndSDOps", - "LoraStateDictWithStrength", - "ModelBuilderProtocol", - "ModuleOps", - "Registry", - "SDKeyValueOperation", - "SDOps", - "SafetensorsModelStateDictLoader", - "SafetensorsStateDictLoader", - "SingleGPUModelBuilder", - "StateDict", - "StateDictLoader", - "StateDictRegistry", - "apply_loras", - "create_meta_model", - "load_state_dict", - "read_model_config", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/attention_ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/attention_ops.py deleted file mode 100644 index 83d244a66..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/attention_ops.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Builder ops for swapping attention backends on a meta model before load.""" - -import torch -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.model.transformer.attention import ( - Attention, - AttentionCallable, - AttentionFunction, - MaskedAttentionCallable, - MaskedAttentionFunction, -) - - -def set_attention_module_op( - attention: AttentionFunction | AttentionCallable | None = None, - masked_attention: MaskedAttentionFunction | MaskedAttentionCallable | None = None, -) -> ModuleOps: - """Build a ``ModuleOps`` that overrides the attention callables on every - ``Attention`` submodule of a model. Applied via ``create_meta_model`` so - the meta model is mutated before weight loading. Matcher returns False - for models with no ``Attention`` submodules, so the op is a no-op there. - Either or both slots may be supplied; *None* leaves that slot untouched. - """ - fn = attention.to_callable() if isinstance(attention, AttentionFunction) else attention - masked_fn = masked_attention.to_callable() if isinstance(masked_attention, MaskedAttentionFunction) else masked_attention - - def matcher(model: torch.nn.Module) -> bool: - return any(isinstance(m, Attention) for m in model.modules()) - - def mutator(model: torch.nn.Module) -> torch.nn.Module: - for module in model.modules(): - if isinstance(module, Attention): - if fn is not None: - module.attention_function = fn - if masked_fn is not None: - module.masked_attention_function = masked_fn - return model - - return ModuleOps(name="set_attention_backend", matcher=matcher, mutator=mutator) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/fuse_loras.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/fuse_loras.py deleted file mode 100644 index 76b9a0a6a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/fuse_loras.py +++ /dev/null @@ -1,177 +0,0 @@ -from collections.abc import Callable, Iterable, Iterator -from dataclasses import dataclass -from typing import NamedTuple - -import torch -from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict - - -class LoraProduct(NamedTuple): - """A LoRA's ``A``, ``B`` factors and its strength scalar.""" - - a: torch.Tensor - b: torch.Tensor - strength: float - - -#: Signature for a fuse callable used by :class:`FuseRule`. -#: -#: Args: -#: key: The state-dict key being fused (e.g. ``"...layers.0.attn.q.weight"``). -#: weight: The current value at ``key`` from ``model_sd``, on the fusion device. -#: deltas: The pre-aggregated LoRA delta for ``key``, in ``aggregation_dtype``. -#: model_sd: The full state dict, for rules that need companion keys -#: (e.g. an existing ``.weight_scale``). -#: -#: Returns a dict of state-dict keys to overwrite -- at minimum ``{key: new_weight}``, -#: plus any companion keys (e.g. an updated ``.weight_scale``) the policy needs to -#: keep in sync. -FuseFn = Callable[[str, torch.Tensor, torch.Tensor, StateDict], dict[str, torch.Tensor]] - - -@dataclass(frozen=True) -class FuseRule: - """Fuse an aggregated LoRA delta into one weight key. - Each policy supplies its own rule (see ``QuantizationPolicy.fuse_rule``); - ``fuse_lora_weights`` is policy-agnostic boilerplate around it. - Attributes: - aggregation_dtype: Dtype callers must pre-aggregate LoRA deltas in - before invoking the rule. - fuse_fn: Callable that applies the pre-aggregated deltas to the weight - (and any companion keys) and returns a dict of keys to overwrite — - at minimum ``{key: new_weight}``, plus any companion keys (e.g. an - updated ``.weight_scale`` for scaled-FP8 layouts) the policy needs - to keep in sync. - """ - - aggregation_dtype: torch.dtype - fuse_fn: FuseFn - - def __call__( - self, - key: str, - weight: torch.Tensor, - deltas: torch.Tensor, - model_sd: StateDict, - ) -> dict[str, torch.Tensor]: - return self.fuse_fn(key, weight, deltas, model_sd) - - -def _bf16_fuse( - key: str, - weight: torch.Tensor, - deltas: torch.Tensor, - model_sd: StateDict, # noqa: ARG001 -) -> dict[str, torch.Tensor]: - deltas.add_(weight) - return {key: deltas.to(dtype=weight.dtype)} - - -bf16_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_bf16_fuse) - - -def _get_device() -> torch.device: - if torch.cuda.is_available(): - return torch.device("cuda", torch.cuda.current_device()) - return torch.device("cpu") - - -def aggregate_lora_products( - products: Iterable[LoraProduct], - dtype: torch.dtype, -) -> torch.Tensor | None: - """Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items. - The first product materializes a freshly-allocated aggregator via - ``torch.matmul(B * strength, A).to(dtype)`` -- preserving the - ``(B * strength) @ A`` rounding pattern. Subsequent products use - ``addmm_`` to avoid allocating the full intermediate delta. - Returns the aggregator, or ``None`` if ``products`` was empty. - """ - aggregated: torch.Tensor | None = None - for product in products: - if aggregated is None: - aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype) - else: - aggregated.addmm_(product.b, product.a, alpha=product.strength) - return aggregated - - -def fuse_lora_weights( - model_sd: StateDict, - lora_sd_and_strengths: list[LoraStateDictWithStrength], - fuse_rule: FuseRule = bf16_fuse_rule, - preserve_input_device: bool = True, -) -> Iterator[tuple[str, torch.Tensor]]: - """Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA. - The fusion math is delegated to ``fuse_rule``. - Output dtypes are the rule's responsibility. - When ``preserve_input_device`` is False, fused tensors are yielded on the device - used for fusion; caller is responsible for moving them to their final - destination. - """ - fusion_device = _get_device() - for key in _affected_weight_keys(lora_sd_and_strengths): - original_weight = model_sd.sd.get(key) - if original_weight is None: - continue - - products = _products_for_sd_key(lora_sd_and_strengths, key, fuse_rule.aggregation_dtype, fusion_device) - deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype) - if deltas is None: - continue - - original_device = original_weight.device - weight = original_weight.to(device=fusion_device) - - fused = fuse_rule(key, weight, deltas, model_sd) - - for k, v in fused.items(): - yield k, v.to(device=original_device) if preserve_input_device else v - - -def apply_loras( - model_sd: StateDict, - lora_sd_and_strengths: list[LoraStateDictWithStrength], - fuse_rule: FuseRule = bf16_fuse_rule, - destination_sd: StateDict | None = None, -) -> StateDict: - """Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``. - When ``destination_sd`` is provided, the fused tensors are placed directly into it. - """ - fused_iter = fuse_lora_weights( - model_sd, - lora_sd_and_strengths, - fuse_rule=fuse_rule, - ) - if destination_sd is not None: - for key, fused in fused_iter: - destination_sd.sd[key] = fused - return destination_sd - - fused = dict(fused_iter) - sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()} - return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype) - - -def _affected_weight_keys(lora_sd_and_strengths: list[LoraStateDictWithStrength]) -> set[str]: - """Return the set of ``.weight`` keys touched by at least one LoRA in the list.""" - suffix = ".lora_A.weight" - return {k[: -len(suffix)] + ".weight" for lsd, _ in lora_sd_and_strengths for k in lsd.sd if k.endswith(suffix)} - - -def _products_for_sd_key( - lora_sd_and_strengths: list[LoraStateDictWithStrength], - key: str, - dtype: torch.dtype, - device: torch.device, -) -> Iterator[LoraProduct]: - """Yield :class:`LoraProduct` items matching *key* across state-dict-backed LoRAs.""" - prefix = key[: -len(".weight")] - key_a = f"{prefix}.lora_A.weight" - key_b = f"{prefix}.lora_B.weight" - for lsd, coef in lora_sd_and_strengths: - if key_a not in lsd.sd or key_b not in lsd.sd: - continue - a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True) - b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True) - yield LoraProduct(a, b, coef) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/helpers.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/helpers.py deleted file mode 100644 index 5119fa15a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/helpers.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Shared model-construction helpers used by both SingleGPUModelBuilder and StreamingModelBuilder.""" - -from __future__ import annotations - -from typing import TypeVar - -import torch -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import StateDict, StateDictLoader -from ltx_core.loader.registry import Registry -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.model_protocol import ModelConfigurator -from torch import nn - -_M = TypeVar("_M", bound=nn.Module) - - -def load_state_dict( - paths: str | tuple[str, ...] | list[str], - loader: StateDictLoader, - registry: Registry, - device: torch.device | None, - sd_ops: SDOps | None = None, -) -> StateDict: - """Load a state dict from disk, using registry caching.""" - if isinstance(paths, str): - path_list = [paths] - elif isinstance(paths, tuple): - path_list = list(paths) - else: - path_list = paths - cached = registry.get(path_list, sd_ops) - if cached is not None: - return cached - result = loader.load(path_list, sd_ops=sd_ops, device=device) - registry.add(path_list, sd_ops=sd_ops, state_dict=result) - return result - - -def read_model_config( - model_path: str | tuple[str, ...], - loader: StateDictLoader, -) -> dict: - """Read metadata from the first shard of a checkpoint.""" - first = model_path[0] if isinstance(model_path, tuple) else model_path - return loader.metadata(first) - - -def create_meta_model( - configurator: type[ModelConfigurator[_M]], - config: dict, - module_ops: tuple[ModuleOps, ...] = (), -) -> _M: - """Create a model on the meta device and apply module operations.""" - with torch.device("meta"): - model = configurator.from_config(config) - for op in module_ops: - if op.matcher(model): - model = op.mutator(model) - return model diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/kernels.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/kernels.py deleted file mode 100644 index 145cd885f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/kernels.py +++ /dev/null @@ -1,79 +0,0 @@ -# ruff: noqa: ANN001, ANN201, ERA001, N803, N806 -try: - import triton - import triton.language as tl - - TRITON_AVAILABLE = True -except (ImportError, OSError): - TRITON_AVAILABLE = False - - -if TRITON_AVAILABLE: - - @triton.jit - def fused_add_round_kernel( - x_ptr, - output_ptr, # contents will be added to the output - seed, - n_elements, - EXPONENT_BIAS, - MANTISSA_BITS, - BLOCK_SIZE: tl.constexpr, - ): - """ - A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding - and add them to bfloat16 output weights. Might be used to upcast original model weights - and to further add them to precalculated deltas coming from LoRAs. - """ - # Get program ID and compute offsets - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - # Load data - x = tl.load(x_ptr + offsets, mask=mask) - rand_vals = tl.rand(seed, offsets) - 0.5 - - x = tl.cast(x, tl.float16) - delta = tl.load(output_ptr + offsets, mask=mask) - delta = tl.cast(delta, tl.float16) - x = x + delta - - x_bits = tl.cast(x, tl.int16, bitcast=True) - - # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for - # normal numbers and -14 for subnormals. - fp16_exponent_bits = (x_bits & 0x7C00) >> 10 - fp16_normals = fp16_exponent_bits > 0 - fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14) - - # Add the target dtype's exponent bias and clamp to the target dtype's exponent range. - exponent = fp16_exponent + EXPONENT_BIAS - MAX_EXPONENT = 2 * EXPONENT_BIAS + 1 - exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent) - exponent = tl.where(exponent < 0, 0, exponent) - - # Normal ULP exponent, expressed as an fp16 exponent field: - # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15 - # Simplifies to: fp16_exponent - MANTISSA_BITS + 15 - # See https://en.wikipedia.org/wiki/Unit_in_the_last_place - eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15)) - - # Calculate epsilon in the target dtype - eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True) - - # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) -> - # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 = - # 16 - EXPONENT_BIAS - MANTISSA_BITS - eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True) - eps = tl.where(exponent > 0, eps_normal, eps_subnormal) - - # Apply zero mask to epsilon - eps = tl.where(x == 0, 0.0, eps) - - # Apply stochastic rounding - output = tl.cast(x + rand_vals * eps, tl.bfloat16) - - # Store the result - tl.store(output_ptr + offsets, output, mask=mask) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/module_ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/module_ops.py deleted file mode 100644 index 0c3ed2f2b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/module_ops.py +++ /dev/null @@ -1,14 +0,0 @@ -from typing import Callable, NamedTuple - -import torch - - -class ModuleOps(NamedTuple): - """ - Defines a named operation for matching and mutating PyTorch modules. - Used to selectively transform modules in a model (e.g., replacing layers with quantized versions). - """ - - name: str - matcher: Callable[[torch.nn.Module], bool] - mutator: Callable[[torch.nn.Module], torch.nn.Module] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/primitives.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/primitives.py deleted file mode 100644 index 4f6b5f8d5..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/primitives.py +++ /dev/null @@ -1,164 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeVar - -import torch -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.sd_ops import SDOps - -if TYPE_CHECKING: - from ltx_core.loader.fuse_loras import FuseRule - from ltx_core.loader.registry import Registry - from typing_extensions import Self - -BuiltType = TypeVar("BuiltType", covariant=True) # noqa: PLC0105 - - -# Per-key shape and dtype description for a flat collection of tensors. -TensorLayout = dict[str, tuple[torch.Size, torch.dtype]] - - -@dataclass(frozen=True) -class StateDict: - """ - Immutable container for a PyTorch state dictionary. - Contains: - - sd: Dictionary of tensors (weights, buffers, etc.) - - device: Device where tensors are stored - - size: Total memory footprint in bytes - - dtype: Set of tensor dtypes present - """ - - sd: dict - device: torch.device - size: int - dtype: set[torch.dtype] - - def footprint(self) -> tuple[int, torch.device]: - return self.size, self.device - - -class StateDictLoader(Protocol): - """ - Protocol for loading state dictionaries from various sources. - Implementations must provide: - - metadata: Extract model metadata from a single path - - load: Load state dict from path(s) and apply SDOps transformations - """ - - def metadata(self, path: str) -> dict: - """ - Load metadata from path - """ - ... - - def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict: - """ - Load state dict from path or paths (for sharded model storage) and apply sd_ops - """ - ... - - -class BuilderProtocol(Protocol[BuiltType]): - """Protocol for model builders that produce a model via ``build()``.""" - - def build( - self, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - **kwargs: Any, # noqa: ANN401 - ) -> BuiltType: ... - - @property - def registry(self) -> "Registry": ... - - def with_registry(self, registry: "Registry") -> "Self": - """Return a copy of this builder using the given weight registry for allocation.""" - ... - - -class ModelBuilderProtocol(BuilderProtocol[BuiltType], Protocol[BuiltType]): - """ - Protocol for building PyTorch models from configuration dictionaries. - Implementations must provide: - - build: Create and initialize a model from state dictionary and apply dtype transformations - """ - - @property - def model_sd_ops(self) -> SDOps | None: ... - - @property - def module_ops(self) -> tuple[ModuleOps, ...]: ... - - @property - def loras(self) -> tuple["LoraPathStrengthAndSDOps", ...]: ... - - def with_sd_ops(self, sd_ops: SDOps | None) -> "Self": - """Return a copy of this builder with the given state-dict key remapping ops.""" - ... - - def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "Self": - """Return a copy of this builder with the given module operations (e.g. quantization).""" - ... - - def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "Self": - """Return a copy of this builder with the given LoRAs to fuse at build time.""" - ... - - def with_lora_load_device(self, device: torch.device) -> "Self": - """Return a copy of this builder that loads LoRA weights onto the given device.""" - ... - - def with_fuse_rule(self, fuse_rule: "FuseRule") -> "Self": - """Return a copy of this builder with the given LoRA fuse rule (e.g. from a quantization policy).""" - ... - - def build( - self, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - **kwargs: Any, # noqa: ANN401 - ) -> BuiltType: - """ - Build the model - Args: - device: Target device for the model - dtype: Target dtype for the model, if None, uses the dtype of the model_path model - Returns: - Model instance - """ - ... - - def model_config(self) -> dict: - """Return the model configuration dictionary extracted from the checkpoint metadata.""" - ... - - -class LoRAAdaptableProtocol(Protocol): - """ - Protocol for models that can be adapted with LoRAs. - Implementations must provide: - - lora: Add a LoRA to the model - """ - - def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "LoRAAdaptableProtocol": ... - - -class LoraPathStrengthAndSDOps(NamedTuple): - """ - Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict. - """ - - path: str - strength: float - sd_ops: SDOps - - -class LoraStateDictWithStrength(NamedTuple): - """ - Tuple containing a LoRA state dict and strength for applying to the model. - """ - - state_dict: StateDict - strength: float diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/registry.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/registry.py deleted file mode 100644 index c49edd581..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/registry.py +++ /dev/null @@ -1,84 +0,0 @@ -import hashlib -import threading -from dataclasses import dataclass, field -from pathlib import Path -from typing import Protocol - -from ltx_core.loader.primitives import StateDict -from ltx_core.loader.sd_ops import SDOps - - -class Registry(Protocol): - """ - Protocol for managing state dictionaries in a registry. - It is used to store state dictionaries and reuse them later without loading them again. - Implementations must provide: - - add: Add a state dictionary to the registry - - pop: Remove a state dictionary from the registry - - get: Retrieve a state dictionary from the registry - - clear: Clear all state dictionaries from the registry - """ - - def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ... - - def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ... - - def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ... - - def clear(self) -> None: ... - - -class DummyRegistry(Registry): - """ - Dummy registry that does not store state dictionaries. - """ - - def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: - pass - - def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: - pass - - def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: - pass - - def clear(self) -> None: - pass - - -@dataclass -class StateDictRegistry(Registry): - """ - Registry that stores state dictionaries in a dictionary. - """ - - _state_dicts: dict[str, StateDict] = field(default_factory=dict) - _lock: threading.Lock = field(default_factory=threading.Lock) - - def _generate_id(self, paths: list[str], sd_ops: SDOps | None) -> str: - m = hashlib.sha256() - parts = [str(Path(p).resolve()) for p in paths] - if sd_ops is not None: - parts.append(sd_ops.name) - m.update("\0".join(parts).encode("utf-8")) - return m.hexdigest() - - def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str: - sd_id = self._generate_id(paths, sd_ops) - with self._lock: - if sd_id in self._state_dicts: - raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first") - self._state_dicts[sd_id] = state_dict - return sd_id - - def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: - with self._lock: - return self._state_dicts.pop(self._generate_id(paths, sd_ops), None) - - def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: - with self._lock: - return self._state_dicts.get(self._generate_id(paths, sd_ops), None) - - def clear(self) -> None: - with self._lock: - self._state_dicts.clear() diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sd_ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sd_ops.py deleted file mode 100644 index 2c5ee23c7..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sd_ops.py +++ /dev/null @@ -1,132 +0,0 @@ -from dataclasses import dataclass, replace -from typing import NamedTuple, Protocol - -import torch - - -@dataclass(frozen=True, slots=True) -class ContentReplacement: - """ - Represents a content replacement operation. - Used to replace a specific content with a replacement in a state dict key. - """ - - content: str - replacement: str - - -@dataclass(frozen=True, slots=True) -class ContentMatching: - """ - Represents a content matching operation. - Used to match a specific prefix and suffix in a state dict key. - """ - - prefix: str = "" - suffix: str = "" - contains: str = "" - - -class KeyValueOperationResult(NamedTuple): - """ - Represents the result of a key-value operation. - Contains the new key and value after the operation has been applied. - """ - - new_key: str - new_value: torch.Tensor - - -class KeyValueOperation(Protocol): - """ - Protocol for key-value operations. - Used to apply operations to a specific key and value in a state dict. - """ - - def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ... - - -@dataclass(frozen=True, slots=True) -class SDKeyValueOperation: - """ - Represents a key-value operation. - Used to apply operations to a specific key and value in a state dict. - """ - - key_matcher: ContentMatching - kv_operation: KeyValueOperation - - -@dataclass(frozen=True, slots=True) -class SDOps: - """Immutable class representing state dict key operations.""" - - name: str - mapping: tuple[ContentReplacement | ContentMatching | SDKeyValueOperation, ...] = () # Immutable tuple of (key, value) pairs - allowed_keys: frozenset[str] | None = None - - def with_replacement(self, content: str, replacement: str) -> "SDOps": - """Create a new SDOps instance with the specified replacement added to the mapping.""" - - new_mapping = (*self.mapping, ContentReplacement(content, replacement)) - return replace(self, mapping=new_mapping) - - def with_matching(self, prefix: str = "", suffix: str = "", contains: str = "") -> "SDOps": - """Create a new SDOps instance with the specified prefix, suffix and contains matching added to the mapping.""" - - new_mapping = (*self.mapping, ContentMatching(prefix, suffix, contains)) - return replace(self, mapping=new_mapping) - - def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps": - """Create a new SDOps instance that only passes keys present in *keys* (post-replacement). - If allowed_keys already exists, the sets are merged via union. - """ - merged = frozenset(keys) | self.allowed_keys if self.allowed_keys is not None else frozenset(keys) - return replace(self, allowed_keys=merged) - - def with_kv_operation( - self, - operation: KeyValueOperation, - key_prefix: str = "", - key_suffix: str = "", - ) -> "SDOps": - """Create a new SDOps instance with the specified value operation added to the mapping.""" - key_matcher = ContentMatching(key_prefix, key_suffix) - sd_kv_operation = SDKeyValueOperation(key_matcher, operation) - new_mapping = (*self.mapping, sd_kv_operation) - return replace(self, mapping=new_mapping) - - def apply_to_key(self, key: str) -> str | None: - """Apply the mapping to the given name.""" - matchers = [content for content in self.mapping if isinstance(content, ContentMatching)] - valid = any(key.startswith(matcher.prefix) and key.endswith(matcher.suffix) and (not matcher.contains or matcher.contains in key) for matcher in matchers) - if not valid: - return None - - for replacement in self.mapping: - if not isinstance(replacement, ContentReplacement): - continue - if replacement.content in key: - key = key.replace(replacement.content, replacement.replacement) - - if self.allowed_keys is not None and key not in self.allowed_keys: - return None - - return key - - def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: - """Apply the value operation to the given name and associated value.""" - for operation in self.mapping: - if not isinstance(operation, SDKeyValueOperation): - continue - if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix): - return operation.kv_operation(key, value) - return [KeyValueOperationResult(key, value)] - - -# Predefined SDOps instances -LTXV_LORA_COMFY_RENAMING_MAP = SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "") - -LTXV_LORA_COMFY_TARGET_MAP = ( - SDOps("LTXV_LORA_COMFY_TARGET_MAP").with_matching().with_replacement("diffusion_model.", "").with_replacement(".lora_A.weight", ".weight").with_replacement(".lora_B.weight", ".weight") -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sft_loader.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sft_loader.py deleted file mode 100644 index 727500a12..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sft_loader.py +++ /dev/null @@ -1,65 +0,0 @@ -import json - -import safetensors -import torch -from ltx_core.loader.primitives import StateDict, StateDictLoader -from ltx_core.loader.sd_ops import SDOps - - -class SafetensorsStateDictLoader(StateDictLoader): - """ - Loads weights from safetensors files without metadata support. - Use this for loading raw weight files. For model files that include - configuration metadata, use SafetensorsModelStateDictLoader instead. - """ - - def metadata(self, path: str) -> dict: - raise NotImplementedError("Not implemented") - - def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict: - """ - Load state dict from path or paths (for sharded model storage) and apply sd_ops - """ - sd = {} - size = 0 - dtype = set() - device = device or torch.device("cpu") - model_paths = path if isinstance(path, list) else [path] - for shard_path in model_paths: - with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f: - safetensor_keys = f.keys() - for name in safetensor_keys: - expected_name = name if sd_ops is None else sd_ops.apply_to_key(name) - if expected_name is None: - continue - value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False) - key_value_pairs = ((expected_name, value),) - if sd_ops is not None: - key_value_pairs = sd_ops.apply_to_key_value(expected_name, value) - for key, value in key_value_pairs: - size += value.nbytes - dtype.add(value.dtype) - sd[key] = value - - return StateDict(sd=sd, device=device, size=size, dtype=dtype) - - -class SafetensorsModelStateDictLoader(StateDictLoader): - """ - Loads weights and configuration metadata from safetensors model files. - Unlike SafetensorsStateDictLoader, this loader can read model configuration - from the safetensors file metadata via the metadata() method. - """ - - def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None): - self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader() - - def metadata(self, path: str) -> dict: - with safetensors.safe_open(path, framework="pt") as f: - meta = f.metadata() - if meta is None or "config" not in meta: - return {} - return json.loads(meta["config"]) - - def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict: - return self.weight_loader.load(path, sd_ops, device) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/single_gpu_model_builder.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/single_gpu_model_builder.py deleted file mode 100644 index f609b8db7..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/single_gpu_model_builder.py +++ /dev/null @@ -1,231 +0,0 @@ -from __future__ import annotations - -import copy -import logging -from typing import TYPE_CHECKING, Final, Generic - -import torch -from ltx_core.loader.fuse_loras import FuseRule, apply_loras, bf16_fuse_rule -from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import ( - LoRAAdaptableProtocol, - LoraPathStrengthAndSDOps, - LoraStateDictWithStrength, - ModelBuilderProtocol, - StateDict, - StateDictLoader, -) -from ltx_core.loader.registry import DummyRegistry, Registry -from ltx_core.loader.sd_ops import SDOps -from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader -from ltx_core.model.model_protocol import ModelConfigurator, ModelType -from torch import nn - -if TYPE_CHECKING: - from typing_extensions import Self - -logger: logging.Logger = logging.getLogger(__name__) - - -def _check_uninitialized(model: nn.Module) -> list[str]: - """Return names of any parameters/buffers still on meta device.""" - names = [] - for name, param in model.named_parameters(): - if str(param.device) == "meta": - names.append(name) - for name, buf in model.named_buffers(): - if str(buf.device) == "meta": - names.append(name) - return names - - -def _load_model_weights( - meta_model: nn.Module, - model_path: str | tuple[str, ...], - loras: tuple[LoraPathStrengthAndSDOps, ...], - loader: StateDictLoader, - registry: Registry, - device: torch.device, - dtype: torch.dtype | None, - model_sd_ops: SDOps | None = None, - lora_load_device: torch.device | None = None, - fuse_rule: FuseRule = bf16_fuse_rule, -) -> None: - """Load base weights and fuse LoRAs into *meta_model* in-place.""" - if lora_load_device is None: - lora_load_device = device - - model_sd = load_state_dict(model_path, loader, registry, device, model_sd_ops) - - lora_strengths = [lora.strength for lora in loras] - if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0): - sd = model_sd.sd - if dtype is not None: - sd = {key: value.to(dtype=dtype) for key, value in sd.items()} - meta_model.load_state_dict(sd, strict=False, assign=True) - return - - lora_state_dicts = [load_state_dict([lora.path], loader, registry, lora_load_device, lora.sd_ops) for lora in loras] - lora_sd_and_strengths = [LoraStateDictWithStrength(sd, strength) for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)] - final_sd = apply_loras( - model_sd=model_sd, - lora_sd_and_strengths=lora_sd_and_strengths, - fuse_rule=fuse_rule, - destination_sd=model_sd if isinstance(registry, DummyRegistry) else None, - ) - fused_sd = final_sd.sd - if dtype is not None: - fused_sd = {key: value.to(dtype=dtype) for key, value in fused_sd.items()} - meta_model.load_state_dict(fused_sd, strict=False, assign=True) - - -class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol): - """ - Builder for PyTorch models residing on a single GPU. - The builder is immutable: ``with_*``/``lora`` return modified copies. The - ``ModelBuilderProtocol`` surface is exposed via read-only properties backed - by private attributes. - Attributes: - model_class_configurator: Class responsible for constructing the model from a config dict. - model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s). - model_sd_ops: Optional state-dict operations applied when loading the model weights. - module_ops: Sequence of module-level mutations applied to the meta model before weight loading. - loras: Sequence of LoRA adapters (path, strength, optional sd_ops) to fuse into the model. - model_loader: Strategy for loading state dicts from disk. Defaults to - :class:`SafetensorsModelStateDictLoader`. - registry: Cache for already-loaded state dicts. Defaults to :class:`DummyRegistry` (no caching). - lora_load_device: Device used when loading LoRA weight tensors from disk. Defaults to - ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to - the target GPU sequentially during fusion, reducing peak GPU memory usage compared to - loading all LoRA weights directly onto the GPU at once. - fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``; - """ - - def __init__( - self, - model_class_configurator: type[ModelConfigurator[ModelType]], - model_path: str | tuple[str, ...], - model_sd_ops: SDOps | None = None, - module_ops: tuple[ModuleOps, ...] = (), - loras: tuple[LoraPathStrengthAndSDOps, ...] = (), - model_loader: StateDictLoader | None = None, - registry: Registry | None = None, - lora_load_device: torch.device | None = None, - fuse_rule: FuseRule = bf16_fuse_rule, - ) -> None: - # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute. - self._model_class_configurator: Final = model_class_configurator - self._model_path = model_path - self._model_sd_ops = model_sd_ops - self._module_ops = module_ops - self._loras = loras - self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader() - self._registry = registry if registry is not None else DummyRegistry() - self._lora_load_device = lora_load_device if lora_load_device is not None else torch.device("cpu") - self._fuse_rule = fuse_rule - - @property - def model_sd_ops(self) -> SDOps | None: - return self._model_sd_ops - - @property - def module_ops(self) -> tuple[ModuleOps, ...]: - return self._module_ops - - @property - def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]: - return self._loras - - @property - def registry(self) -> Registry: - return self._registry - - @property - def model_path(self) -> str | tuple[str, ...]: - return self._model_path - - @property - def model_loader(self) -> StateDictLoader: - return self._model_loader - - @property - def lora_load_device(self) -> torch.device: - return self._lora_load_device - - @property - def fuse_rule(self) -> FuseRule: - return self._fuse_rule - - def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> Self: - clone = copy.copy(self) - clone._loras = (*self._loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)) - return clone - - def with_sd_ops(self, sd_ops: SDOps | None) -> Self: - clone = copy.copy(self) - clone._model_sd_ops = sd_ops - return clone - - def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self: - clone = copy.copy(self) - clone._module_ops = module_ops - return clone - - def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self: - clone = copy.copy(self) - clone._loras = loras - return clone - - def with_registry(self, registry: Registry) -> Self: - clone = copy.copy(self) - clone._registry = registry - return clone - - def with_lora_load_device(self, device: torch.device) -> Self: - clone = copy.copy(self) - clone._lora_load_device = device - return clone - - def with_fuse_rule(self, fuse_rule: FuseRule) -> Self: - clone = copy.copy(self) - clone._fuse_rule = fuse_rule - return clone - - def model_config(self) -> dict: - return read_model_config(self._model_path, self._model_loader) - - def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType: - return create_meta_model(self._model_class_configurator, config, module_ops) - - def load_sd(self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None) -> StateDict: - return load_state_dict(paths, self._model_loader, registry, device, sd_ops) - - def build( - self, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - **kwargs: object, # noqa: ARG002 - ) -> ModelType: - device = torch.device("cuda") if device is None else device - config = self.model_config() - meta_model = self.meta_model(config, self._module_ops) - - _load_model_weights( - meta_model=meta_model, - model_path=self._model_path, - loras=self._loras, - loader=self._model_loader, - registry=self._registry, - device=device, - dtype=dtype, - model_sd_ops=self._model_sd_ops, - lora_load_device=self._lora_load_device, - fuse_rule=self._fuse_rule, - ) - - uninitialized = _check_uninitialized(meta_model) - if uninitialized: - logger.warning(f"Uninitialized parameters or buffers: {uninitialized}") - return meta_model - return meta_model.to(device) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/modality_tiling.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/modality_tiling.py deleted file mode 100644 index 0f1b445a1..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/modality_tiling.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Video modality tiling helpers. -Provides :class:`VideoModalityTilingHelper` — a stateless helper that -tiles and blends video :class:`Modality` token sequences by -spatial/temporal region. Tile geometry is represented by the existing -:class:`Tile` NamedTuple from :mod:`ltx_core.tiling`; no distributed -primitives are required. -""" - -from __future__ import annotations - -from dataclasses import dataclass, replace - -import torch -from ltx_core.model.transformer.modality import Modality -from ltx_core.tiling import Tile, TileCountConfig, create_tiles, identity_mapping_operation, split_by_count -from ltx_core.tools import VideoLatentTools -from ltx_core.types import VideoLatentShape - - -@dataclass(frozen=True) -class TilingContext: - """Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`. - Carries the token-level keep indices and per-conditioning-token blend - weights needed by :meth:`~VideoModalityTilingHelper.blend`. - """ - - keep_indices: torch.Tensor - """``(num_kept,)`` int64 — sorted indices of tokens the tile processes.""" - num_total_tokens: int - """Total number of tokens in the full (untiled) sequence.""" - cond_blend_weights: torch.Tensor | None - """``(num_kept_cond,)`` — weight for each kept conditioning token, - equal to ``1 / num_tiles_that_keep_this_token``. ``None`` when - there are no conditioning tokens.""" - - -class VideoModalityTilingHelper: - """Stateless helper that tiles and blends video :class:`Modality` sequences. - Constructed once with a :class:`TileCountConfig` and - :class:`VideoLatentTools`. Tiles are computed at construction and - available via the :attr:`tiles` property. Use :meth:`tile_modality` - and :meth:`blend` with any tile from that list. - Usage:: - helper = VideoModalityTilingHelper(tiling, video_tools) - for tile in helper.tiles: - tiled_mod, ctx = helper.tile_modality(modality, tile) - result = run_model(tiled_mod) - helper.blend(result, tile, ctx, output=output) - """ - - def __init__(self, tiling: TileCountConfig, video_tools: VideoLatentTools) -> None: - self._patchifier = video_tools.patchifier - self._latent_shape = video_tools.target_shape - self._num_generated_tokens = self._patchifier.get_token_count(self._latent_shape) - self._tiles = create_tiles( - torch.Size([self._latent_shape.frames, self._latent_shape.height, self._latent_shape.width]), - splitters=[ - split_by_count(tiling.frames.num_tiles, tiling.frames.overlap), - split_by_count(tiling.height.num_tiles, tiling.height.overlap), - split_by_count(tiling.width.num_tiles, tiling.width.overlap), - ], - mappers=[identity_mapping_operation] * 3, - ) - - @property - def tiles(self) -> list[Tile]: - """All tiles for the configured tiling layout.""" - return self._tiles - - # -- tile modality ----------------------------------------------------- - - def tile_modality(self, modality: Modality, tile: Tile, *, normalize_positions: bool = True) -> tuple[Modality, TilingContext]: - """Slice *modality* to the tokens covered by *tile*. - Selects generated tokens belonging to the tile's spatial region - and conditioning tokens that overlap with the tile (or have - negative time coordinates). - Args: - normalize_positions: When True, shift all positions so the - tile's generated tokens start at zero in every dimension. - Returns: - A ``(tiled_modality, context)`` tuple. Pass *context* to - :meth:`blend` together with the model output. - """ - device = modality.positions.device - gen_indices = self._generated_token_indices(tile, device=device) - num_total = modality.latent.shape[1] - - cond_blend_weights: torch.Tensor | None = None - if num_total > self._num_generated_tokens: - keep_per_tile_cond = self._all_tiles_cond_keep(modality) # (num_tiles, num_cond) bool - tile_idx = next((i for i, t in enumerate(self._tiles) if t.in_coords == tile.in_coords), None) - if tile_idx is None: - raise ValueError(f"Tile with in_coords={tile.in_coords} is not in this helper's tile set; pass a tile obtained from `helper.tiles`.") - my_cond_keep = keep_per_tile_cond[tile_idx] - cond_indices = self._num_generated_tokens + my_cond_keep.nonzero(as_tuple=False).squeeze(1) - keep_indices = torch.cat([gen_indices, cond_indices]) - total_keepers = keep_per_tile_cond.sum(dim=0).float() # (num_cond,) - cond_blend_weights = 1.0 / total_keepers[my_cond_keep] - else: - keep_indices = gen_indices - - tile_attention_mask = None - if modality.attention_mask is not None: - tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices] - - positions = modality.positions[:, :, keep_indices, :] - if normalize_positions: - num_tile_gen = self._tile_generated_token_count(tile) - gen_pos = positions[:, :, :num_tile_gen, :] # (B, 3, num_tile_gen, 2) - offset = gen_pos[..., 0].amin(dim=2, keepdim=True).unsqueeze(-1) # (B, 3, 1, 1) - positions = positions - offset - - tiled = replace( - modality, - latent=modality.latent[:, keep_indices, :], - timesteps=modality.timesteps[:, keep_indices], - positions=positions, - attention_mask=tile_attention_mask, - ) - - return tiled, TilingContext(keep_indices=keep_indices, num_total_tokens=num_total, cond_blend_weights=cond_blend_weights) - - # -- blend ------------------------------------------------------------- - - def blend( - self, - tile_to_blend: torch.Tensor, - tile: Tile, - context: TilingContext, - output: torch.Tensor | None = None, - ) -> torch.Tensor: - """Blend-weight tile results and accumulate into the full token space. - Premultiplied (blend-weighted) data is **added** to *output*, - allowing multiple tiles to be accumulated into the same buffer. - Args: - tile_to_blend: Denoised tile tensor ``(B, num_tile_tokens, D)``, - where the first ``_tile_generated_token_count(tile)`` - entries are generated tokens and the remainder are - conditioning tokens. - tile: The :class:`Tile` that was used in :meth:`tile_modality`. - context: The :class:`TilingContext` returned by :meth:`tile_modality`. - output: Optional pre-allocated output tensor. When provided - its shape must be ``(B, num_total_tokens, D)`` and the - blended tile is **added** into it. When ``None`` a new - zero-filled tensor is created. - Returns: - The output tensor with the blended tile added at the correct - positions. - """ - batch, _, dim = tile_to_blend.shape - num_tile_gen = self._tile_generated_token_count(tile) - gen_indices = self._generated_token_indices(tile, device=tile_to_blend.device) - - num_total_tokens = context.num_total_tokens - expected_shape = (batch, num_total_tokens, dim) - - if output is not None: - if output.shape != expected_shape: - raise ValueError(f"Expected output shape {expected_shape}, got {output.shape}") - result = output - else: - result = torch.zeros(*expected_shape, device=tile_to_blend.device, dtype=tile_to_blend.dtype) - - # Blend mask is (tile_F, tile_H, tile_W) — one weight per token in row-major order. - blend_weights = tile.blend_mask.reshape(-1).to(device=tile_to_blend.device, dtype=tile_to_blend.dtype) - tile_gen = tile_to_blend[:, :num_tile_gen, :] * blend_weights[None, :, None] - - result[:, gen_indices, :] += tile_gen - - # Scatter kept conditioning tokens, weighted by 1/N where N is - # the number of tiles that keep each token (so they sum to 1). - if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None: - cond_indices = context.keep_indices[context.keep_indices >= self._num_generated_tokens] - weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype) - result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None] - - return result - - # -- private ----------------------------------------------------------- - - def _tile_generated_token_count(self, tile: Tile) -> int: - """Number of generated tokens in *tile*.""" - frame_slice, height_slice, width_slice = tile.in_coords - tile_shape = VideoLatentShape( - batch=self._latent_shape.batch, - channels=self._latent_shape.channels, - frames=frame_slice.stop - frame_slice.start, - height=height_slice.stop - height_slice.start, - width=width_slice.stop - width_slice.start, - ) - return self._patchifier.get_token_count(tile_shape) - - def _generated_token_indices(self, tile: Tile, device: torch.device | None = None) -> torch.Tensor: - """Flat token indices of *tile*'s generated tokens in the full sequence.""" - frame_slice, height_slice, width_slice = tile.in_coords - f = torch.arange(frame_slice.start, frame_slice.stop, device=device) - h = torch.arange(height_slice.start, height_slice.stop, device=device) - w = torch.arange(width_slice.start, width_slice.stop, device=device) - return (f[:, None, None] * self._latent_shape.height * self._latent_shape.width + h[None, :, None] * self._latent_shape.width + w[None, None, :]).reshape(-1) - - def _all_tiles_cond_keep(self, modality: Modality) -> torch.Tensor: - """Vectorized (num_tiles, num_cond) bool: which tiles keep each conditioning token. - A conditioning token is kept by a tile when its ``[start, end)`` interval - overlaps the tile in all three dimensions, or when it has a negative time - coordinate (reference token). - """ - cond_positions = modality.positions[:, :, self._num_generated_tokens :, :] # (B, 3, num_cond, 2) - device = cond_positions.device - - # Per-tile (start, end) bounds along each axis; small Python loop (num_tiles <= ~16). - starts_list: list[torch.Tensor] = [] - ends_list: list[torch.Tensor] = [] - for t in self._tiles: - gen_idx = self._generated_token_indices(t, device=device) - gen_positions = modality.positions[:, :, gen_idx, :] # (B, 3, num_tile_gen, 2) - starts_list.append(gen_positions[..., 0].amin(dim=2)) # (B, 3) - ends_list.append(gen_positions[..., 1].amax(dim=2)) # (B, 3) - tile_starts = torch.stack(starts_list, dim=0) # (num_tiles, B, 3) - tile_ends = torch.stack(ends_list, dim=0) # (num_tiles, B, 3) - - cond_starts = cond_positions[..., 0] # (B, 3, num_cond) - cond_ends = cond_positions[..., 1] # (B, 3, num_cond) - # Broadcast: (1, B, 3, num_cond) vs (num_tiles, B, 3, 1) -> (num_tiles, B, 3, num_cond). - overlaps = (cond_starts[None] < tile_ends[..., None]) & (cond_ends[None] > tile_starts[..., None]) - overlaps_all_dims = overlaps.all(dim=2) # (num_tiles, B, num_cond) - has_negative_time = (cond_positions[:, 0, :, 0] < 0)[None] # (1, B, num_cond) - return (overlaps_all_dims | has_negative_time).any(dim=1) # (num_tiles, num_cond) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/__init__.py deleted file mode 100644 index e8a5b27ab..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Model definitions for LTX-2.""" - -from ltx_core.model.model_protocol import ModelConfigurator, ModelType - -__all__ = [ - "ModelConfigurator", - "ModelType", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/__init__.py deleted file mode 100644 index c1cd3ac31..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Audio VAE model components.""" - -from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio, encode_audio -from ltx_core.model.audio_vae.model_configurator import ( - AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, - AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, - VOCODER_COMFY_KEYS_FILTER, - AudioDecoderConfigurator, - AudioEncoderConfigurator, - VocoderConfigurator, -) -from ltx_core.model.audio_vae.ops import AudioProcessor -from ltx_core.model.audio_vae.vocoder import Vocoder, VocoderWithBWE - -__all__ = [ - "AUDIO_VAE_DECODER_COMFY_KEYS_FILTER", - "AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER", - "VOCODER_COMFY_KEYS_FILTER", - "AudioDecoder", - "AudioDecoderConfigurator", - "AudioEncoder", - "AudioEncoderConfigurator", - "AudioProcessor", - "Vocoder", - "VocoderConfigurator", - "VocoderWithBWE", - "decode_audio", - "encode_audio", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/attention.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/attention.py deleted file mode 100644 index 5e3adef70..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/attention.py +++ /dev/null @@ -1,70 +0,0 @@ -from enum import Enum - -import torch -from ltx_core.model.common.normalization import NormType, build_normalization_layer - - -class AttentionType(Enum): - """Enum for specifying the attention mechanism type.""" - - VANILLA = "vanilla" - LINEAR = "linear" - NONE = "none" - - -class AttnBlock(torch.nn.Module): - def __init__( - self, - in_channels: int, - norm_type: NormType = NormType.GROUP, - ) -> None: - super().__init__() - self.in_channels = in_channels - - self.norm = build_normalization_layer(in_channels, normtype=norm_type) - self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - h_ = x - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - # compute attention - b, c, h, w = q.shape - q = q.reshape(b, c, h * w).contiguous() - q = q.permute(0, 2, 1).contiguous() # b,hw,c - k = k.reshape(b, c, h * w).contiguous() # b,c,hw - w_ = torch.bmm(q, k).contiguous() # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] - w_ = w_ * (int(c) ** (-0.5)) - w_ = torch.nn.functional.softmax(w_, dim=2) - - # attend to values - v = v.reshape(b, c, h * w).contiguous() - w_ = w_.permute(0, 2, 1).contiguous() # b,hw,hw (first hw of k, second of q) - h_ = torch.bmm(v, w_).contiguous() # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] - h_ = h_.reshape(b, c, h, w).contiguous() - - h_ = self.proj_out(h_) - - return x + h_ - - -def make_attn( - in_channels: int, - attn_type: AttentionType = AttentionType.VANILLA, - norm_type: NormType = NormType.GROUP, -) -> torch.nn.Module: - match attn_type: - case AttentionType.VANILLA: - return AttnBlock(in_channels, norm_type=norm_type) - case AttentionType.NONE: - return torch.nn.Identity() - case AttentionType.LINEAR: - raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.") - case _: - raise ValueError(f"Unknown attention type: {attn_type}") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/audio_vae.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/audio_vae.py deleted file mode 100644 index ad8e48b99..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/audio_vae.py +++ /dev/null @@ -1,501 +0,0 @@ -from typing import Set, Tuple - -import torch -import torch.nn.functional as F -from ltx_core.components.patchifiers import AudioPatchifier -from ltx_core.model.audio_vae.attention import AttentionType, make_attn -from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d -from ltx_core.model.audio_vae.causality_axis import CausalityAxis -from ltx_core.model.audio_vae.downsample import build_downsampling_path -from ltx_core.model.audio_vae.ops import AudioProcessor, PerChannelStatistics -from ltx_core.model.audio_vae.resnet import ResnetBlock -from ltx_core.model.audio_vae.upsample import build_upsampling_path -from ltx_core.model.audio_vae.vocoder import Vocoder -from ltx_core.model.common.normalization import NormType, build_normalization_layer -from ltx_core.types import Audio, AudioLatentShape - -LATENT_DOWNSAMPLE_FACTOR = 4 - - -def build_mid_block( - channels: int, - temb_channels: int, - dropout: float, - norm_type: NormType, - causality_axis: CausalityAxis, - attn_type: AttentionType, - add_attention: bool, -) -> torch.nn.Module: - """Build the middle block with two ResNet blocks and optional attention.""" - mid = torch.nn.Module() - mid.block_1 = ResnetBlock( - in_channels=channels, - out_channels=channels, - temb_channels=temb_channels, - dropout=dropout, - norm_type=norm_type, - causality_axis=causality_axis, - ) - mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity() - mid.block_2 = ResnetBlock( - in_channels=channels, - out_channels=channels, - temb_channels=temb_channels, - dropout=dropout, - norm_type=norm_type, - causality_axis=causality_axis, - ) - return mid - - -def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor: - """Run features through the middle block.""" - features = mid.block_1(features, temb=None) - features = mid.attn_1(features) - return mid.block_2(features, temb=None) - - -class AudioEncoder(torch.nn.Module): - """ - Encoder that compresses audio spectrograms into latent representations. - The encoder uses a series of downsampling blocks with residual connections, - attention mechanisms, and configurable causal convolutions. - """ - - def __init__( # noqa: PLR0913 - self, - *, - ch: int, - ch_mult: Tuple[int, ...] = (1, 2, 4, 8), - num_res_blocks: int, - attn_resolutions: Set[int], - dropout: float = 0.0, - resamp_with_conv: bool = True, - in_channels: int, - resolution: int, - z_channels: int, - double_z: bool = True, - attn_type: AttentionType = AttentionType.VANILLA, - mid_block_add_attention: bool = True, - norm_type: NormType = NormType.GROUP, - causality_axis: CausalityAxis = CausalityAxis.WIDTH, - sample_rate: int = 16000, - mel_hop_length: int = 160, - n_fft: int = 1024, - is_causal: bool = True, - mel_bins: int = 64, - **_ignore_kwargs, - ) -> None: - """ - Initialize the Encoder. - Args: - Arguments are configuration parameters, loaded from the audio VAE checkpoint config - (audio_vae.model.params.ddconfig): - ch: Base number of feature channels used in the first convolution layer. - ch_mult: Multiplicative factors for the number of channels at each resolution level. - num_res_blocks: Number of residual blocks to use at each resolution level. - attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention. - resolution: Input spatial resolution of the spectrogram (height, width). - z_channels: Number of channels in the latent representation. - norm_type: Normalization layer type to use within the network (e.g., group, batch). - causality_axis: Axis along which convolutions should be causal (e.g., time axis). - sample_rate: Audio sample rate in Hz for the input signals. - mel_hop_length: Hop length used when computing the mel spectrogram. - n_fft: FFT size used to compute the spectrogram. - mel_bins: Number of mel-frequency bins in the input spectrogram. - in_channels: Number of channels in the input spectrogram tensor. - double_z: If True, predict both mean and log-variance (doubling latent channels). - is_causal: If True, use causal convolutions suitable for streaming setups. - dropout: Dropout probability used in residual and mid blocks. - attn_type: Type of attention mechanism to use in attention blocks. - resamp_with_conv: If True, perform resolution changes using strided convolutions. - mid_block_add_attention: If True, add an attention block in the mid-level of the encoder. - """ - super().__init__() - - self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) - self.sample_rate = sample_rate - self.mel_hop_length = mel_hop_length - self.n_fft = n_fft - self.is_causal = is_causal - self.mel_bins = mel_bins - - self.patchifier = AudioPatchifier( - patch_size=1, - audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, - sample_rate=sample_rate, - hop_length=mel_hop_length, - is_causal=is_causal, - ) - - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.z_channels = z_channels - self.double_z = double_z - self.norm_type = norm_type - self.causality_axis = causality_axis - self.attn_type = attn_type - - # downsampling - self.conv_in = make_conv2d( - in_channels, - self.ch, - kernel_size=3, - stride=1, - causality_axis=self.causality_axis, - ) - - self.non_linearity = torch.nn.SiLU() - - self.down, block_in = build_downsampling_path( - ch=ch, - ch_mult=ch_mult, - num_resolutions=self.num_resolutions, - num_res_blocks=num_res_blocks, - resolution=resolution, - temb_channels=self.temb_ch, - dropout=dropout, - norm_type=self.norm_type, - causality_axis=self.causality_axis, - attn_type=self.attn_type, - attn_resolutions=attn_resolutions, - resamp_with_conv=resamp_with_conv, - ) - - self.mid = build_mid_block( - channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - norm_type=self.norm_type, - causality_axis=self.causality_axis, - attn_type=self.attn_type, - add_attention=mid_block_add_attention, - ) - - self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type) - self.conv_out = make_conv2d( - block_in, - 2 * z_channels if double_z else z_channels, - kernel_size=3, - stride=1, - causality_axis=self.causality_axis, - ) - - def forward(self, spectrogram: torch.Tensor) -> torch.Tensor: - """ - Encode audio spectrogram into latent representations. - Args: - spectrogram: Input spectrogram of shape (batch, channels, time, frequency) - Returns: - Encoded latent representation of shape (batch, channels, frames, mel_bins) - """ - h = self.conv_in(spectrogram) - h = self._run_downsampling_path(h) - h = run_mid_block(self.mid, h) - h = self._finalize_output(h) - - return self._normalize_latents(h) - - def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor: - for level in range(self.num_resolutions): - stage = self.down[level] - for block_idx in range(self.num_res_blocks): - h = stage.block[block_idx](h, temb=None) - if stage.attn: - h = stage.attn[block_idx](h) - - if level != self.num_resolutions - 1: - h = stage.downsample(h) - - return h - - def _finalize_output(self, h: torch.Tensor) -> torch.Tensor: - h = self.norm_out(h) - h = self.non_linearity(h) - return self.conv_out(h) - - def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor: - """ - Normalize encoder latents using per-channel statistics. - When the encoder is configured with ``double_z=True``, the final - convolution produces twice the number of latent channels, typically - interpreted as two concatenated tensors along the channel dimension - (e.g., mean and variance or other auxiliary parameters). - This method intentionally uses only the first half of the channels - (the "mean" component) as input to the patchifier and normalization - logic. The remaining channels are left unchanged by this method and - are expected to be consumed elsewhere in the VAE pipeline. - If ``double_z=False``, the encoder output already contains only the - mean latents and the chunking operation simply returns that tensor. - """ - means = torch.chunk(latent_output, 2, dim=1)[0] - latent_shape = AudioLatentShape( - batch=means.shape[0], - channels=means.shape[1], - frames=means.shape[2], - mel_bins=means.shape[3], - ) - latent_patched = self.patchifier.patchify(means) - latent_normalized = self.per_channel_statistics.normalize(latent_patched) - return self.patchifier.unpatchify(latent_normalized, latent_shape) - - -def encode_audio( - audio: Audio, - audio_encoder: AudioEncoder, - audio_processor: AudioProcessor | None = None, -) -> torch.Tensor: - """Encode audio waveform into latent representation. - Args: - audio: Audio container with waveform tensor of shape (batch, channels, samples) and sampling rate. - audio_encoder: Audio encoder model - audio_processor: Audio processor model (optional, if not provided, it will be created from the audio encoder) - """ - dtype = next(audio_encoder.parameters()).dtype - device = next(audio_encoder.parameters()).device - - if audio_processor is None: - audio_processor = AudioProcessor( - target_sample_rate=audio_encoder.sample_rate, - mel_bins=audio_encoder.mel_bins, - mel_hop_length=audio_encoder.mel_hop_length, - n_fft=audio_encoder.n_fft, - ).to(device=device) - - mel_spectrogram = audio_processor.waveform_to_mel(audio.to(device=device)) - - latent = audio_encoder(mel_spectrogram.to(dtype=dtype)) - return latent - - -class AudioDecoder(torch.nn.Module): - """ - Symmetric decoder that reconstructs audio spectrograms from latent features. - The decoder mirrors the encoder structure with configurable channel multipliers, - attention resolutions, and causal convolutions. - """ - - def __init__( # noqa: PLR0913 - self, - *, - ch: int, - out_ch: int, - ch_mult: Tuple[int, ...] = (1, 2, 4, 8), - num_res_blocks: int, - attn_resolutions: Set[int], - resolution: int, - z_channels: int, - norm_type: NormType = NormType.GROUP, - causality_axis: CausalityAxis = CausalityAxis.WIDTH, - dropout: float = 0.0, - mid_block_add_attention: bool = True, - sample_rate: int = 16000, - mel_hop_length: int = 160, - is_causal: bool = True, - mel_bins: int | None = None, - ) -> None: - """ - Initialize the Decoder. - Args: - Arguments are configuration parameters, loaded from the audio VAE checkpoint config - (audio_vae.model.params.ddconfig): - - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions - - resolution, z_channels - - norm_type, causality_axis - """ - super().__init__() - - # Internal behavioural defaults that are not driven by the checkpoint. - resamp_with_conv = True - attn_type = AttentionType.VANILLA - - # Per-channel statistics for denormalizing latents - self.per_channel_statistics = PerChannelStatistics(latent_channels=ch) - self.sample_rate = sample_rate - self.mel_hop_length = mel_hop_length - self.is_causal = is_causal - self.mel_bins = mel_bins - self.patchifier = AudioPatchifier( - patch_size=1, - audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR, - sample_rate=sample_rate, - hop_length=mel_hop_length, - is_causal=is_causal, - ) - - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.out_ch = out_ch - self.give_pre_end = False - self.tanh_out = False - self.norm_type = norm_type - self.z_channels = z_channels - self.channel_multipliers = ch_mult - self.attn_resolutions = attn_resolutions - self.causality_axis = causality_axis - self.attn_type = attn_type - - base_block_channels = ch * self.channel_multipliers[-1] - base_resolution = resolution // (2 ** (self.num_resolutions - 1)) - self.z_shape = (1, z_channels, base_resolution, base_resolution) - - self.conv_in = make_conv2d(z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis) - self.non_linearity = torch.nn.SiLU() - self.mid = build_mid_block( - channels=base_block_channels, - temb_channels=self.temb_ch, - dropout=dropout, - norm_type=self.norm_type, - causality_axis=self.causality_axis, - attn_type=self.attn_type, - add_attention=mid_block_add_attention, - ) - self.up, final_block_channels = build_upsampling_path( - ch=ch, - ch_mult=ch_mult, - num_resolutions=self.num_resolutions, - num_res_blocks=num_res_blocks, - resolution=resolution, - temb_channels=self.temb_ch, - dropout=dropout, - norm_type=self.norm_type, - causality_axis=self.causality_axis, - attn_type=self.attn_type, - attn_resolutions=attn_resolutions, - resamp_with_conv=resamp_with_conv, - initial_block_channels=base_block_channels, - ) - - self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type) - self.conv_out = make_conv2d(final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis) - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - """ - Decode latent features back to audio spectrograms. - Args: - sample: Encoded latent representation of shape (batch, channels, frames, mel_bins) - Returns: - Reconstructed audio spectrogram of shape (batch, channels, time, frequency) - """ - sample, target_shape = self._denormalize_latents(sample) - - h = self.conv_in(sample) - h = run_mid_block(self.mid, h) - h = self._run_upsampling_path(h) - h = self._finalize_output(h) - - return self._adjust_output_shape(h, target_shape) - - def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]: - latent_shape = AudioLatentShape( - batch=sample.shape[0], - channels=sample.shape[1], - frames=sample.shape[2], - mel_bins=sample.shape[3], - ) - - sample_patched = self.patchifier.patchify(sample) - sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched) - sample = self.patchifier.unpatchify(sample_denormalized, latent_shape) - - target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR - if self.causality_axis != CausalityAxis.NONE: - target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1) - - target_shape = AudioLatentShape( - batch=latent_shape.batch, - channels=self.out_ch, - frames=target_frames, - mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins, - ) - - return sample, target_shape - - def _adjust_output_shape( - self, - decoded_output: torch.Tensor, - target_shape: AudioLatentShape, - ) -> torch.Tensor: - """ - Adjust output shape to match target dimensions for variable-length audio. - This function handles the common case where decoded audio spectrograms need to be - resized to match a specific target shape. - Args: - decoded_output: Tensor of shape (batch, channels, time, frequency) - target_shape: AudioLatentShape describing (batch, channels, time, mel bins) - Returns: - Tensor adjusted to match target_shape exactly - """ - # Current output shape: (batch, channels, time, frequency) - _, _, current_time, current_freq = decoded_output.shape - target_channels = target_shape.channels - target_time = target_shape.frames - target_freq = target_shape.mel_bins - - # Step 1: Crop first to avoid exceeding target dimensions - decoded_output = decoded_output[:, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)] - - # Step 2: Calculate padding needed for time and frequency dimensions - time_padding_needed = target_time - decoded_output.shape[2] - freq_padding_needed = target_freq - decoded_output.shape[3] - - # Step 3: Apply padding if needed - if time_padding_needed > 0 or freq_padding_needed > 0: - # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom) - # For audio: pad_left/right = frequency, pad_top/bottom = time - padding = ( - 0, - max(freq_padding_needed, 0), # frequency padding (left, right) - 0, - max(time_padding_needed, 0), # time padding (top, bottom) - ) - decoded_output = F.pad(decoded_output, padding) - - # Step 4: Final safety crop to ensure exact target shape - decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq] - - return decoded_output - - def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor: - for level in reversed(range(self.num_resolutions)): - stage = self.up[level] - for block_idx, block in enumerate(stage.block): - h = block(h, temb=None) - if stage.attn: - h = stage.attn[block_idx](h) - - if level != 0 and hasattr(stage, "upsample"): - h = stage.upsample(h) - - return h - - def _finalize_output(self, h: torch.Tensor) -> torch.Tensor: - if self.give_pre_end: - return h - - h = self.norm_out(h) - h = self.non_linearity(h) - h = self.conv_out(h) - return torch.tanh(h) if self.tanh_out else h - - -def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> Audio: - """ - Decode an audio latent representation using the provided audio decoder and vocoder. - Args: - latent: Input audio latent tensor. - audio_decoder: Model to decode the latent to waveform features. - vocoder: Model to convert decoded features to audio waveform. - Returns: - Decoded audio with waveform and sampling rate. - """ - decoded_audio = audio_decoder(latent) - waveform = vocoder(decoded_audio).squeeze(0).float() - return Audio(waveform=waveform, sampling_rate=vocoder.output_sampling_rate) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causal_conv_2d.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causal_conv_2d.py deleted file mode 100644 index 6a9dd3cba..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causal_conv_2d.py +++ /dev/null @@ -1,109 +0,0 @@ -import torch -import torch.nn.functional as F -from ltx_core.model.audio_vae.causality_axis import CausalityAxis - - -class CausalConv2d(torch.nn.Module): - """ - A causal 2D convolution. - This layer ensures that the output at time `t` only depends on inputs - at time `t` and earlier. It achieves this by applying asymmetric padding - to the time dimension (width) before the convolution. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int | tuple[int, int], - stride: int = 1, - dilation: int | tuple[int, int] = 1, - groups: int = 1, - bias: bool = True, - causality_axis: CausalityAxis = CausalityAxis.HEIGHT, - ) -> None: - super().__init__() - - self.causality_axis = causality_axis - - # Ensure kernel_size and dilation are tuples - kernel_size = torch.nn.modules.utils._pair(kernel_size) - dilation = torch.nn.modules.utils._pair(dilation) - - # Calculate padding dimensions - pad_h = (kernel_size[0] - 1) * dilation[0] - pad_w = (kernel_size[1] - 1) * dilation[1] - - # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom) - match self.causality_axis: - case CausalityAxis.NONE: - self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) - case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY: - self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) - case CausalityAxis.HEIGHT: - self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) - case _: - raise ValueError(f"Invalid causality_axis: {causality_axis}") - - # The internal convolution layer uses no padding, as we handle it manually - self.conv = torch.nn.Conv2d( - in_channels, - out_channels, - kernel_size, - stride=stride, - padding=0, - dilation=dilation, - groups=groups, - bias=bias, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # Apply causal padding before convolution - x = F.pad(x, self.padding) - return self.conv(x) - - -def make_conv2d( - in_channels: int, - out_channels: int, - kernel_size: int | tuple[int, int], - stride: int = 1, - padding: tuple[int, int, int, int] | None = None, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - causality_axis: CausalityAxis | None = None, -) -> torch.nn.Module: - """ - Create a 2D convolution layer that can be either causal or non-causal. - Args: - in_channels: Number of input channels - out_channels: Number of output channels - kernel_size: Size of the convolution kernel - stride: Convolution stride - padding: Padding (if None, will be calculated based on causal flag) - dilation: Dilation rate - groups: Number of groups for grouped convolution - bias: Whether to use bias - causality_axis: Dimension along which to apply causality. - Returns: - Either a regular Conv2d or CausalConv2d layer - """ - if causality_axis is not None: - # For causal convolution, padding is handled internally by CausalConv2d - return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis) - else: - # For non-causal convolution, use symmetric padding if not specified - if padding is None: - padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size) - - return torch.nn.Conv2d( - in_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - groups, - bias, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causality_axis.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causality_axis.py deleted file mode 100644 index b99f83550..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causality_axis.py +++ /dev/null @@ -1,10 +0,0 @@ -from enum import Enum - - -class CausalityAxis(Enum): - """Enum for specifying the causality axis in causal convolutions.""" - - NONE = None - WIDTH = "width" - HEIGHT = "height" - WIDTH_COMPATIBILITY = "width-compatibility" diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/downsample.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/downsample.py deleted file mode 100644 index fb01a015d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/downsample.py +++ /dev/null @@ -1,109 +0,0 @@ -from typing import Set, Tuple - -import torch -from ltx_core.model.audio_vae.attention import AttentionType, make_attn -from ltx_core.model.audio_vae.causality_axis import CausalityAxis -from ltx_core.model.audio_vae.resnet import ResnetBlock -from ltx_core.model.common.normalization import NormType - - -class Downsample(torch.nn.Module): - """ - A downsampling layer that can use either a strided convolution - or average pooling. Supports standard and causal padding for the - convolutional mode. - """ - - def __init__( - self, - in_channels: int, - with_conv: bool, - causality_axis: CausalityAxis = CausalityAxis.WIDTH, - ) -> None: - super().__init__() - self.with_conv = with_conv - self.causality_axis = causality_axis - - if self.causality_axis != CausalityAxis.NONE and not self.with_conv: - raise ValueError("causality is only supported when `with_conv=True`.") - - if self.with_conv: - # Do time downsampling here - # no asymmetric padding in torch conv, must do it ourselves - self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.with_conv: - # Padding tuple is in the order: (left, right, top, bottom). - match self.causality_axis: - case CausalityAxis.NONE: - pad = (0, 1, 0, 1) - case CausalityAxis.WIDTH: - pad = (2, 0, 0, 1) - case CausalityAxis.HEIGHT: - pad = (0, 1, 2, 0) - case CausalityAxis.WIDTH_COMPATIBILITY: - pad = (1, 0, 0, 1) - case _: - raise ValueError(f"Invalid causality_axis: {self.causality_axis}") - - x = torch.nn.functional.pad(x, pad, mode="constant", value=0) - x = self.conv(x) - else: - # This branch is only taken if with_conv=False, which implies causality_axis is NONE. - x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) - - return x - - -def build_downsampling_path( # noqa: PLR0913 - *, - ch: int, - ch_mult: Tuple[int, ...], - num_resolutions: int, - num_res_blocks: int, - resolution: int, - temb_channels: int, - dropout: float, - norm_type: NormType, - causality_axis: CausalityAxis, - attn_type: AttentionType, - attn_resolutions: Set[int], - resamp_with_conv: bool, -) -> tuple[torch.nn.ModuleList, int]: - """Build the downsampling path with residual blocks, attention, and downsampling layers.""" - down_modules = torch.nn.ModuleList() - curr_res = resolution - in_ch_mult = (1, *tuple(ch_mult)) - block_in = ch - - for i_level in range(num_resolutions): - block = torch.nn.ModuleList() - attn = torch.nn.ModuleList() - block_in = ch * in_ch_mult[i_level] - block_out = ch * ch_mult[i_level] - - for _ in range(num_res_blocks): - block.append( - ResnetBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=temb_channels, - dropout=dropout, - norm_type=norm_type, - causality_axis=causality_axis, - ) - ) - block_in = block_out - if curr_res in attn_resolutions: - attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) - - down = torch.nn.Module() - down.block = block - down.attn = attn - if i_level != num_resolutions - 1: - down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis) - curr_res = curr_res // 2 - down_modules.append(down) - - return down_modules, block_in diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/model_configurator.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/model_configurator.py deleted file mode 100644 index 359637596..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/model_configurator.py +++ /dev/null @@ -1,193 +0,0 @@ -import torch -from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps -from ltx_core.model.audio_vae.attention import AttentionType -from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder -from ltx_core.model.audio_vae.causality_axis import CausalityAxis -from ltx_core.model.audio_vae.vocoder import MelSTFT, Vocoder, VocoderWithBWE -from ltx_core.model.common.normalization import NormType -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.utils import check_config_value - - -def _vocoder_from_config( - cfg: dict, - apply_final_activation: bool = True, - output_sampling_rate: int | None = None, -) -> Vocoder: - """Instantiate a Vocoder from a flat config dict. - Args: - cfg: Vocoder config dict (keys match Vocoder constructor args). - apply_final_activation: Whether to apply tanh/clamp at the output. - output_sampling_rate: Explicit override for the output sample rate. - When None, reads from cfg["output_sampling_rate"] (default 24000). - """ - return Vocoder( - resblock_kernel_sizes=cfg.get("resblock_kernel_sizes", [3, 7, 11]), - upsample_rates=cfg.get("upsample_rates", [6, 5, 2, 2, 2]), - upsample_kernel_sizes=cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]), - resblock_dilation_sizes=cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]), - upsample_initial_channel=cfg.get("upsample_initial_channel", 1024), - resblock=cfg.get("resblock", "1"), - output_sampling_rate=(output_sampling_rate if output_sampling_rate is not None else cfg.get("output_sampling_rate", 24000)), - activation=cfg.get("activation", "snake"), - use_tanh_at_final=cfg.get("use_tanh_at_final", True), - apply_final_activation=apply_final_activation, - use_bias_at_final=cfg.get("use_bias_at_final", True), - ) - - -class VocoderConfigurator(ModelConfigurator[Vocoder]): - """Configurator that auto-detects the checkpoint format. - Returns a plain Vocoder for pre-ltx-2.3 checkpoints (flat config) or a - VocoderWithBWE for ltx-2.3+ checkpoints (nested "vocoder" + "bwe" config). - """ - - @classmethod - def from_config(cls: type[Vocoder], config: dict) -> Vocoder | VocoderWithBWE: - cfg = config.get("vocoder", {}) - - if "bwe" not in cfg: - check_config_value(cfg, "resblock", "1") - check_config_value(cfg, "stereo", True) - return _vocoder_from_config(cfg) - - vocoder_cfg = cfg.get("vocoder", {}) - bwe_cfg = cfg["bwe"] - - check_config_value(vocoder_cfg, "resblock", "AMP1") - check_config_value(vocoder_cfg, "stereo", True) - check_config_value(vocoder_cfg, "activation", "snakebeta") - check_config_value(bwe_cfg, "resblock", "AMP1") - check_config_value(bwe_cfg, "stereo", True) - check_config_value(bwe_cfg, "activation", "snakebeta") - - vocoder = _vocoder_from_config( - vocoder_cfg, - output_sampling_rate=bwe_cfg["input_sampling_rate"], - ) - bwe_generator = _vocoder_from_config( - bwe_cfg, - apply_final_activation=False, - output_sampling_rate=bwe_cfg["output_sampling_rate"], - ) - mel_stft = MelSTFT( - filter_length=bwe_cfg["n_fft"], - hop_length=bwe_cfg["hop_length"], - win_length=bwe_cfg["n_fft"], - n_mel_channels=bwe_cfg["num_mels"], - ) - return VocoderWithBWE( - vocoder=vocoder, - bwe_generator=bwe_generator, - mel_stft=mel_stft, - input_sampling_rate=bwe_cfg["input_sampling_rate"], - output_sampling_rate=bwe_cfg["output_sampling_rate"], - hop_length=bwe_cfg["hop_length"], - ) - - -def _strip_vocoder_prefix(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: - """Strip the leading 'vocoder.' prefix exactly once. - Uses removeprefix instead of str.replace so that BWE keys like - 'vocoder.vocoder.conv_pre' become 'vocoder.conv_pre' (not 'conv_pre'). - Works identically for legacy keys like 'vocoder.conv_pre' → 'conv_pre'. - """ - return [KeyValueOperationResult(key.removeprefix("vocoder."), value)] - - -VOCODER_COMFY_KEYS_FILTER = SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_kv_operation(operation=_strip_vocoder_prefix, key_prefix="vocoder.") - - -class AudioDecoderConfigurator(ModelConfigurator[AudioDecoder]): - @classmethod - def from_config(cls: type[AudioDecoder], config: dict) -> AudioDecoder: - audio_vae_cfg = config.get("audio_vae", {}) - model_cfg = audio_vae_cfg.get("model", {}) - model_params = model_cfg.get("params", {}) - ddconfig = model_params.get("ddconfig", {}) - preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) - stft_cfg = preprocessing_cfg.get("stft", {}) - mel_cfg = preprocessing_cfg.get("mel", {}) - variables_cfg = audio_vae_cfg.get("variables", {}) - - sample_rate = model_params.get("sampling_rate", 16000) - mel_hop_length = stft_cfg.get("hop_length", 160) - is_causal = stft_cfg.get("causal", True) - mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") - - return AudioDecoder( - ch=ddconfig.get("ch", 128), - out_ch=ddconfig.get("out_ch", 2), - ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), - num_res_blocks=ddconfig.get("num_res_blocks", 2), - attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}), - resolution=ddconfig.get("resolution", 256), - z_channels=ddconfig.get("z_channels", 8), - norm_type=NormType(ddconfig.get("norm_type", "pixel")), - causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), - dropout=ddconfig.get("dropout", 0.0), - mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), - sample_rate=sample_rate, - mel_hop_length=mel_hop_length, - is_causal=is_causal, - mel_bins=mel_bins, - ) - - -class AudioEncoderConfigurator(ModelConfigurator[AudioEncoder]): - @classmethod - def from_config(cls: type[AudioEncoder], config: dict) -> AudioEncoder: - audio_vae_cfg = config.get("audio_vae", {}) - model_cfg = audio_vae_cfg.get("model", {}) - model_params = model_cfg.get("params", {}) - ddconfig = model_params.get("ddconfig", {}) - preprocessing_cfg = audio_vae_cfg.get("preprocessing", {}) - stft_cfg = preprocessing_cfg.get("stft", {}) - mel_cfg = preprocessing_cfg.get("mel", {}) - variables_cfg = audio_vae_cfg.get("variables", {}) - - sample_rate = model_params.get("sampling_rate", 16000) - mel_hop_length = stft_cfg.get("hop_length", 160) - n_fft = stft_cfg.get("filter_length", 1024) - is_causal = stft_cfg.get("causal", True) - mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins") - - return AudioEncoder( - ch=ddconfig.get("ch", 128), - ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))), - num_res_blocks=ddconfig.get("num_res_blocks", 2), - attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}), - resolution=ddconfig.get("resolution", 256), - z_channels=ddconfig.get("z_channels", 8), - double_z=ddconfig.get("double_z", True), - dropout=ddconfig.get("dropout", 0.0), - resamp_with_conv=ddconfig.get("resamp_with_conv", True), - in_channels=ddconfig.get("in_channels", 2), - attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")), - mid_block_add_attention=ddconfig.get("mid_block_add_attention", True), - norm_type=NormType(ddconfig.get("norm_type", "pixel")), - causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")), - sample_rate=sample_rate, - mel_hop_length=mel_hop_length, - n_fft=n_fft, - is_causal=is_causal, - mel_bins=mel_bins, - ) - - -AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = ( - SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER") - .with_matching(prefix="audio_vae.decoder.") - .with_matching(prefix="audio_vae.per_channel_statistics.") - .with_replacement("audio_vae.decoder.", "") - .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.") -) - - -AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = ( - SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER") - .with_matching(prefix="audio_vae.encoder.") - .with_matching(prefix="audio_vae.per_channel_statistics.") - .with_replacement("audio_vae.encoder.", "") - .with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.") -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/ops.py deleted file mode 100644 index fa34b7d85..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/ops.py +++ /dev/null @@ -1,72 +0,0 @@ -import torch -import torchaudio -from ltx_core.types import Audio -from torch import nn - - -class AudioProcessor(nn.Module): - """Converts audio waveforms to log-mel spectrograms with optional resampling.""" - - def __init__( - self, - target_sample_rate: int, - mel_bins: int, - mel_hop_length: int, - n_fft: int, - ) -> None: - super().__init__() - self.target_sample_rate = target_sample_rate - self.mel_transform = torchaudio.transforms.MelSpectrogram( - sample_rate=target_sample_rate, - n_fft=n_fft, - win_length=n_fft, - hop_length=mel_hop_length, - f_min=0.0, - f_max=target_sample_rate / 2.0, - n_mels=mel_bins, - window_fn=torch.hann_window, - center=True, - pad_mode="reflect", - power=1.0, - mel_scale="slaney", - norm="slaney", - ) - - def resample_audio(self, audio: Audio) -> Audio: - """Resample audio to the processor's target sample rate if needed.""" - if audio.sampling_rate == self.target_sample_rate: - return audio - resampled = torchaudio.functional.resample(audio.waveform, audio.sampling_rate, self.target_sample_rate) - resampled = resampled.to(device=audio.waveform.device, dtype=audio.waveform.dtype) - return Audio(waveform=resampled, sampling_rate=self.target_sample_rate) - - def waveform_to_mel( - self, - audio: Audio, - ) -> torch.Tensor: - """Convert waveform to log-mel spectrogram [batch, channels, time, n_mels].""" - waveform = self.resample_audio(audio).waveform - - mel = self.mel_transform(waveform) - mel = torch.log(torch.clamp(mel, min=1e-5)) - - mel = mel.to(device=waveform.device, dtype=waveform.dtype) - return mel.permute(0, 1, 3, 2).contiguous() - - -class PerChannelStatistics(nn.Module): - """ - Per-channel statistics for normalizing and denormalizing the latent representation. - This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict. - """ - - def __init__(self, latent_channels: int = 128) -> None: - super().__init__() - self.register_buffer("std-of-means", torch.empty(latent_channels)) - self.register_buffer("mean-of-means", torch.empty(latent_channels)) - - def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x) - - def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/resnet.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/resnet.py deleted file mode 100644 index 37c868139..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/resnet.py +++ /dev/null @@ -1,171 +0,0 @@ -from typing import Tuple - -import torch -from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d -from ltx_core.model.audio_vae.causality_axis import CausalityAxis -from ltx_core.model.common.normalization import NormType, build_normalization_layer - -LRELU_SLOPE = 0.1 - - -class ResBlock1(torch.nn.Module): - def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)): - super(ResBlock1, self).__init__() - self.convs1 = torch.nn.ModuleList( - [ - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - padding="same", - ), - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - padding="same", - ), - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[2], - padding="same", - ), - ] - ) - - self.convs2 = torch.nn.ModuleList( - [ - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding="same", - ), - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding="same", - ), - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=1, - padding="same", - ), - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for conv1, conv2 in zip(self.convs1, self.convs2, strict=True): - xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE) - xt = conv1(xt) - xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE) - xt = conv2(xt) - x = xt + x - return x - - -class ResBlock2(torch.nn.Module): - def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)): - super(ResBlock2, self).__init__() - self.convs = torch.nn.ModuleList( - [ - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - padding="same", - ), - torch.nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - padding="same", - ), - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for conv in self.convs: - xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE) - xt = conv(xt) - x = xt + x - return x - - -class ResnetBlock(torch.nn.Module): - def __init__( - self, - *, - in_channels: int, - out_channels: int | None = None, - conv_shortcut: bool = False, - dropout: float = 0.0, - temb_channels: int = 512, - norm_type: NormType = NormType.GROUP, - causality_axis: CausalityAxis = CausalityAxis.HEIGHT, - ) -> None: - super().__init__() - self.causality_axis = causality_axis - - if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP: - raise ValueError("Causal ResnetBlock with GroupNorm is not supported.") - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - - self.norm1 = build_normalization_layer(in_channels, normtype=norm_type) - self.non_linearity = torch.nn.SiLU() - self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) - if temb_channels > 0: - self.temb_proj = torch.nn.Linear(temb_channels, out_channels) - self.norm2 = build_normalization_layer(out_channels, normtype=norm_type) - self.dropout = torch.nn.Dropout(dropout) - self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis) - else: - self.nin_shortcut = make_conv2d(in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis) - - def forward( - self, - x: torch.Tensor, - temb: torch.Tensor | None = None, - ) -> torch.Tensor: - h = x - h = self.norm1(h) - h = self.non_linearity(h) - h = self.conv1(h) - - if temb is not None: - h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None] - - h = self.norm2(h) - h = self.non_linearity(h) - h = self.dropout(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x) - - return x + h diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/upsample.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/upsample.py deleted file mode 100644 index 6181492d4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/upsample.py +++ /dev/null @@ -1,105 +0,0 @@ -from typing import Set, Tuple - -import torch -from ltx_core.model.audio_vae.attention import AttentionType, make_attn -from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d -from ltx_core.model.audio_vae.causality_axis import CausalityAxis -from ltx_core.model.audio_vae.resnet import ResnetBlock -from ltx_core.model.common.normalization import NormType - - -class Upsample(torch.nn.Module): - def __init__( - self, - in_channels: int, - with_conv: bool, - causality_axis: CausalityAxis = CausalityAxis.HEIGHT, - ) -> None: - super().__init__() - self.with_conv = with_conv - self.causality_axis = causality_axis - if self.with_conv: - self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") - if self.with_conv: - x = self.conv(x) - # Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n. - # For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2]. - # The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2], - # So the output elements rely on the following windows: - # 0: [-,-,0] - # 1: [-,0,0] - # 2: [0,0,1] - # 3: [0,1,1] - # 4: [1,1,2] - # 5: [1,2,2] - # Notice that the first and second elements in the output rely only on the first element in the input, - # while all other elements rely on two elements in the input. - # So we can drop the first element to undo the padding (rather than the last element). - # This is a no-op for non-causal convolutions. - match self.causality_axis: - case CausalityAxis.NONE: - pass # x remains unchanged - case CausalityAxis.HEIGHT: - x = x[:, :, 1:, :] - case CausalityAxis.WIDTH: - x = x[:, :, :, 1:] - case CausalityAxis.WIDTH_COMPATIBILITY: - pass # x remains unchanged - case _: - raise ValueError(f"Invalid causality_axis: {self.causality_axis}") - - return x - - -def build_upsampling_path( # noqa: PLR0913 - *, - ch: int, - ch_mult: Tuple[int, ...], - num_resolutions: int, - num_res_blocks: int, - resolution: int, - temb_channels: int, - dropout: float, - norm_type: NormType, - causality_axis: CausalityAxis, - attn_type: AttentionType, - attn_resolutions: Set[int], - resamp_with_conv: bool, - initial_block_channels: int, -) -> tuple[torch.nn.ModuleList, int]: - """Build the upsampling path with residual blocks, attention, and upsampling layers.""" - up_modules = torch.nn.ModuleList() - block_in = initial_block_channels - curr_res = resolution // (2 ** (num_resolutions - 1)) - - for level in reversed(range(num_resolutions)): - stage = torch.nn.Module() - stage.block = torch.nn.ModuleList() - stage.attn = torch.nn.ModuleList() - block_out = ch * ch_mult[level] - - for _ in range(num_res_blocks + 1): - stage.block.append( - ResnetBlock( - in_channels=block_in, - out_channels=block_out, - temb_channels=temb_channels, - dropout=dropout, - norm_type=norm_type, - causality_axis=causality_axis, - ) - ) - block_in = block_out - if curr_res in attn_resolutions: - stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) - - if level != 0: - stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis) - curr_res *= 2 - - up_modules.insert(0, stage) - - return up_modules, block_in diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/vocoder.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/vocoder.py deleted file mode 100644 index 31ce7d765..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/vocoder.py +++ /dev/null @@ -1,591 +0,0 @@ -import math -from typing import List - -import einops -import torch -import torch.nn.functional as F -from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1 -from torch import nn - - -def get_padding(kernel_size: int, dilation: int = 1) -> int: - return int((kernel_size * dilation - dilation) / 2) - - -# --------------------------------------------------------------------------- -# Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2 -# Adopted from https://github.com/NVIDIA/BigVGAN -# --------------------------------------------------------------------------- - - -def _sinc(x: torch.Tensor) -> torch.Tensor: - return torch.where( - x == 0, - torch.tensor(1.0, device=x.device, dtype=x.dtype), - torch.sin(math.pi * x) / math.pi / x, - ) - - -def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: - even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - delta_f = 4 * half_width - amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - if amplitude > 50.0: - beta = 0.1102 * (amplitude - 8.7) - elif amplitude >= 21.0: - beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0) - else: - beta = 0.0 - window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size - if cutoff == 0: - filter_ = torch.zeros_like(time) - else: - filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) - filter_ /= filter_.sum() - return filter_.view(1, 1, kernel_size) - - -class LowPassFilter1d(nn.Module): - def __init__( - self, - cutoff: float = 0.5, - half_width: float = 0.6, - stride: int = 1, - padding: bool = True, - padding_mode: str = "replicate", - kernel_size: int = 12, - ) -> None: - super().__init__() - if cutoff < -0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - self.kernel_size = kernel_size - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = stride - self.padding = padding - self.padding_mode = padding_mode - self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - _, n_channels, _ = x.shape - if self.padding: - x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) - return F.conv1d(x, self.filter.expand(n_channels, -1, -1), stride=self.stride, groups=n_channels) - - -class UpSample1d(nn.Module): - def __init__( - self, - ratio: int = 2, - kernel_size: int | None = None, - persistent: bool = True, - window_type: str = "kaiser", - ) -> None: - super().__init__() - self.ratio = ratio - self.stride = ratio - - if window_type == "hann": - # Hann-windowed sinc filter equivalent to torchaudio.functional.resample - rolloff = 0.99 - lowpass_filter_width = 6 - width = math.ceil(lowpass_filter_width / rolloff) - self.kernel_size = 2 * width * ratio + 1 - self.pad = width - self.pad_left = 2 * width * ratio - self.pad_right = self.kernel_size - ratio - time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff - time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) - window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 - sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1) - else: - # Kaiser-windowed sinc filter (BigVGAN default). - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - sinc_filter = kaiser_sinc_filter1d( - cutoff=0.5 / ratio, - half_width=0.6 / ratio, - kernel_size=self.kernel_size, - ) - - self.register_buffer("filter", sinc_filter, persistent=persistent) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - _, n_channels, _ = x.shape - x = F.pad(x, (self.pad, self.pad), mode="replicate") - filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) - x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels) - return x[..., self.pad_left : -self.pad_right] - - -class DownSample1d(nn.Module): - def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None: - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.lowpass = LowPassFilter1d( - cutoff=0.5 / ratio, - half_width=0.6 / ratio, - stride=ratio, - kernel_size=self.kernel_size, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.lowpass(x) - - -class Activation1d(nn.Module): - def __init__( - self, - activation: nn.Module, - up_ratio: int = 2, - down_ratio: int = 2, - up_kernel_size: int = 12, - down_kernel_size: int = 12, - ) -> None: - super().__init__() - self.act = activation - self.upsample = UpSample1d(up_ratio, up_kernel_size) - self.downsample = DownSample1d(down_ratio, down_kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.upsample(x) - x = self.act(x) - return self.downsample(x) - - -class Snake(nn.Module): - def __init__( - self, - in_features: int, - alpha: float = 1.0, - alpha_trainable: bool = True, - alpha_logscale: bool = True, - ) -> None: - super().__init__() - self.alpha_logscale = alpha_logscale - self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) - self.alpha.requires_grad = alpha_trainable - self.eps = 1e-9 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2) - - -class SnakeBeta(nn.Module): - def __init__( - self, - in_features: int, - alpha: float = 1.0, - alpha_trainable: bool = True, - alpha_logscale: bool = True, - ) -> None: - super().__init__() - self.alpha_logscale = alpha_logscale - self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) - self.alpha.requires_grad = alpha_trainable - self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) - self.beta.requires_grad = alpha_trainable - self.eps = 1e-9 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) - - -class AMPBlock1(nn.Module): - def __init__( - self, - channels: int, - kernel_size: int = 3, - dilation: tuple[int, int, int] = (1, 3, 5), - activation: str = "snake", - ) -> None: - super().__init__() - act_cls = SnakeBeta if activation == "snakebeta" else Snake - self.convs1 = nn.ModuleList( - [ - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[0], - padding=get_padding(kernel_size, dilation[0]), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[1], - padding=get_padding(kernel_size, dilation[1]), - ), - nn.Conv1d( - channels, - channels, - kernel_size, - 1, - dilation=dilation[2], - padding=get_padding(kernel_size, dilation[2]), - ), - ] - ) - - self.convs2 = nn.ModuleList( - [ - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), - nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), - ] - ) - - self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))]) - self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))]) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True): - xt = a1(x) - xt = c1(xt) - xt = a2(xt) - xt = c2(xt) - x = x + xt - return x - - -class Vocoder(torch.nn.Module): - """ - Vocoder model for synthesizing audio from Mel spectrograms. - Args: - resblock_kernel_sizes: List of kernel sizes for the residual blocks. - This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`. - upsample_rates: List of upsampling rates. - This value is read from the checkpoint at `config.vocoder.upsample_rates`. - upsample_kernel_sizes: List of kernel sizes for the upsampling layers. - This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`. - resblock_dilation_sizes: List of dilation sizes for the residual blocks. - This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`. - upsample_initial_channel: Initial number of channels for the upsampling layers. - This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`. - resblock: Type of residual block to use ("1", "2", or "AMP1"). - This value is read from the checkpoint at `config.vocoder.resblock`. - output_sampling_rate: Waveform sample rate. - This value is read from the checkpoint at `config.vocoder.output_sampling_rate`. - activation: Activation type for BigVGAN v2 ("snake" or "snakebeta"). Only used when resblock="AMP1". - use_tanh_at_final: Apply tanh at the output (when apply_final_activation=True). - apply_final_activation: Whether to apply the final tanh/clamp activation. - use_bias_at_final: Whether to use bias in the final conv layer. - """ - - def __init__( # noqa: PLR0913 - self, - resblock_kernel_sizes: List[int] | None = None, - upsample_rates: List[int] | None = None, - upsample_kernel_sizes: List[int] | None = None, - resblock_dilation_sizes: List[List[int]] | None = None, - upsample_initial_channel: int = 1024, - resblock: str = "1", - output_sampling_rate: int = 24000, - activation: str = "snake", - use_tanh_at_final: bool = True, - apply_final_activation: bool = True, - use_bias_at_final: bool = True, - ) -> None: - super().__init__() - - # Mutable default values are not supported as default arguments. - if resblock_kernel_sizes is None: - resblock_kernel_sizes = [3, 7, 11] - if upsample_rates is None: - upsample_rates = [6, 5, 2, 2, 2] - if upsample_kernel_sizes is None: - upsample_kernel_sizes = [16, 15, 8, 4, 4] - if resblock_dilation_sizes is None: - resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] - - self.output_sampling_rate = output_sampling_rate - self.num_kernels = len(resblock_kernel_sizes) - self.num_upsamples = len(upsample_rates) - self.use_tanh_at_final = use_tanh_at_final - self.apply_final_activation = apply_final_activation - self.is_amp = resblock == "AMP1" - - # All production checkpoints are stereo: 128 input channels (2 stereo channels x 64 mel - # bins each), 2 output channels. - self.conv_pre = nn.Conv1d( - in_channels=128, - out_channels=upsample_initial_channel, - kernel_size=7, - stride=1, - padding=3, - ) - resblock_cls = ResBlock1 if resblock == "1" else AMPBlock1 - - self.ups = nn.ModuleList( - nn.ConvTranspose1d( - upsample_initial_channel // (2**i), - upsample_initial_channel // (2 ** (i + 1)), - kernel_size, - stride, - padding=(kernel_size - stride) // 2, - ) - for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)) - ) - - final_channels = upsample_initial_channel // (2 ** len(upsample_rates)) - self.resblocks = nn.ModuleList() - - for i in range(len(upsample_rates)): - ch = upsample_initial_channel // (2 ** (i + 1)) - for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): - if self.is_amp: - self.resblocks.append(resblock_cls(ch, kernel_size, dilations, activation=activation)) - else: - self.resblocks.append(resblock_cls(ch, kernel_size, dilations)) - - if self.is_amp: - self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels)) - else: - self.act_post = nn.LeakyReLU() - - # All production checkpoints are stereo: this final conv maps `final_channels` to 2 output channels (stereo). - self.conv_post = nn.Conv1d( - in_channels=final_channels, - out_channels=2, - kernel_size=7, - stride=1, - padding=3, - bias=use_bias_at_final, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the vocoder. - Args: - x: Input Mel spectrogram tensor. Can be either: - - 3D: (batch_size, time, mel_bins) for mono - - 4D: (batch_size, 2, time, mel_bins) for stereo - Returns: - Audio waveform tensor of shape (batch_size, out_channels, audio_length) - """ - x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time) - - if x.dim() == 4: # stereo - assert x.shape[1] == 2, "Input must have 2 channels for stereo" - x = einops.rearrange(x, "b s c t -> b (s c) t") - - x = self.conv_pre(x) - - for i in range(self.num_upsamples): - if not self.is_amp: - x = F.leaky_relu(x, LRELU_SLOPE) - x = self.ups[i](x) - start = i * self.num_kernels - end = start + self.num_kernels - - # Evaluate all resblocks with the same input tensor so they can run - # independently (and thus in parallel on accelerator hardware) before - # aggregating their outputs via mean. - block_outputs = torch.stack( - [self.resblocks[idx](x) for idx in range(start, end)], - dim=0, - ) - x = block_outputs.mean(dim=0) - - x = self.act_post(x) - x = self.conv_post(x) - - if self.apply_final_activation: - x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1) - - return x - - -class _STFTFn(nn.Module): - """Implements STFT as a convolution with precomputed DFT x Hann-window bases. - The DFT basis rows (real and imaginary parts interleaved) multiplied by the causal - Hann window are stored as buffers and loaded from the checkpoint. Using the exact - bfloat16 bases from training ensures the mel values fed to the BWE generator are - bit-identical to what it was trained on. - """ - - def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None: - super().__init__() - self.hop_length = hop_length - self.win_length = win_length - n_freqs = filter_length // 2 + 1 - self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) - self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)) - - def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Compute magnitude and phase spectrogram from a batch of waveforms. - Applies causal (left-only) padding of win_length - hop_length samples so that - each output frame depends only on past and present input — no lookahead. - Args: - y: Waveform tensor of shape (B, T). - Returns: - magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames). - phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames). - """ - if y.dim() == 2: - y = y.unsqueeze(1) # (B, 1, T) - left_pad = max(0, self.win_length - self.hop_length) # causal: left-only - y = F.pad(y, (left_pad, 0)) - spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) - n_freqs = spec.shape[1] // 2 - real, imag = spec[:, :n_freqs], spec[:, n_freqs:] - magnitude = torch.sqrt(real**2 + imag**2) - phase = torch.atan2(imag.float(), real.float()).to(real.dtype) - return magnitude, phase - - -class MelSTFT(nn.Module): - """Causal log-mel spectrogram module whose buffers are loaded from the checkpoint. - Computes a log-mel spectrogram by running the causal STFT (_STFTFn) on the input - waveform and projecting the linear magnitude spectrum onto the mel filterbank. - The module's state dict layout matches the 'mel_stft.*' keys stored in the checkpoint - (mel_basis, stft_fn.forward_basis, stft_fn.inverse_basis). - """ - - def __init__( - self, - filter_length: int, - hop_length: int, - win_length: int, - n_mel_channels: int, - ) -> None: - super().__init__() - self.stft_fn = _STFTFn(filter_length, hop_length, win_length) - - # Initialized to zeros; load_state_dict overwrites with the checkpoint's - # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]). - n_freqs = filter_length // 2 + 1 - self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) - - def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute log-mel spectrogram and auxiliary spectral quantities. - Args: - y: Waveform tensor of shape (B, T). - Returns: - log_mel: Log-compressed mel spectrogram, shape (B, n_mel_channels, T_frames). - magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames). - phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames). - energy: Per-frame energy (L2 norm over frequency), shape (B, T_frames). - """ - magnitude, phase = self.stft_fn(y) - energy = torch.norm(magnitude, dim=1) - mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) - log_mel = torch.log(torch.clamp(mel, min=1e-5)) - return log_mel, magnitude, phase, energy - - -class VocoderWithBWE(nn.Module): - """Vocoder with bandwidth extension (BWE) upsampling. - Chains a mel-to-wav vocoder with a BWE module that upsamples the output - to a higher sample rate. The BWE computes a mel spectrogram from the - vocoder output, runs it through a second generator to predict a residual, - and adds it to a sinc-resampled skip connection. - The forward pass runs in fp32 via autocast to avoid bfloat16 accumulation - errors that degrade spectral metrics by 40-90%. - """ - - def __init__( - self, - vocoder: Vocoder, - bwe_generator: Vocoder, - mel_stft: MelSTFT, - input_sampling_rate: int, - output_sampling_rate: int, - hop_length: int, - ) -> None: - super().__init__() - self.vocoder = vocoder - self.bwe_generator = bwe_generator - self.mel_stft = mel_stft - self.input_sampling_rate = input_sampling_rate - self.output_sampling_rate = output_sampling_rate - self.hop_length = hop_length - # Compute the resampler on CPU so the sinc filter is materialized even when - # the model is constructed on meta device (SingleGPUModelBuilder pattern). - # The filter is not stored in the checkpoint (persistent=False). - with torch.device("cpu"): - self.resampler = UpSample1d(ratio=output_sampling_rate // input_sampling_rate, persistent=False, window_type="hann") - - @property - def conv_pre(self) -> nn.Conv1d: - return self.vocoder.conv_pre - - @property - def conv_post(self) -> nn.Conv1d: - return self.vocoder.conv_post - - def _compute_mel(self, audio: torch.Tensor) -> torch.Tensor: - """Compute log-mel spectrogram from waveform using causal STFT bases. - Args: - audio: Waveform tensor of shape (B, C, T). - Returns: - mel: Log-mel spectrogram of shape (B, C, n_mels, T_frames). - """ - batch, n_channels, _ = audio.shape - flat = audio.reshape(batch * n_channels, -1) # (B*C, T) - mel, _, _, _ = self.mel_stft.mel_spectrogram(flat) # (B*C, n_mels, T_frames) - return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2]) # (B, C, n_mels, T_frames) - - def forward(self, mel_spec: torch.Tensor) -> torch.Tensor: - """Run the full vocoder + BWE forward pass. - Runs in float32 regardless of weight or input dtype. bfloat16 arithmetic - causes 40-90% spectral metric degradation due to accumulation errors - compounding through 108 sequential convolutions in the BigVGAN v2 architecture. - Args: - mel_spec: Mel spectrogram of shape (B, 2, T, mel_bins) for stereo - or (B, T, mel_bins) for mono. Same format as Vocoder.forward. - Returns: - Waveform tensor of shape (B, out_channels, T_out) clipped to [-1, 1]. - """ - input_dtype = mel_spec.dtype - # Run the entire forward pass in fp32. bfloat16 accumulation errors - # compound through 108 sequential convolutions and degrade spectral - # metrics (mel_l1, MRSTFT) by 40-90% while perceptual quality (CDPAM) - # is unaffected. fp32 eliminates this degradation. - # We use autocast(dtype=float32) rather than self.float() because it - # upcasts bf16 weights per-op at kernel level, avoiding the temporary - # memory spike of self.float() / self.to(original_dtype). - # Benchmarked on H100 (128.5M-param model): - # autocast fp32: +70 MB peak VRAM, 123 ms (vs 482 MB / 95 ms for bf16) - # model.float(): +324 MB peak VRAM, 149 ms - # Tested: both approaches produce bit-identical output. - - with torch.autocast(device_type=mel_spec.device.type, dtype=torch.float32): - x = self.vocoder(mel_spec.float()) - _, _, length_low_rate = x.shape - output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate - - # Pad to multiple of hop_length for exact mel frame count - remainder = length_low_rate % self.hop_length - if remainder != 0: - x = F.pad(x, (0, self.hop_length - remainder)) - - # Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames) - mel = self._compute_mel(x) - - # Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator - mel_for_bwe = mel.transpose(2, 3) # (B, C, T_frames, mel_bins) - residual = self.bwe_generator(mel_for_bwe) - skip = self.resampler(x) - assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}" - - return torch.clamp(residual + skip, -1, 1)[..., :output_length].to(input_dtype) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/__init__.py deleted file mode 100644 index 42713f7d0..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Common model utilities.""" - -from ltx_core.model.common.normalization import NormType, PixelNorm, build_normalization_layer - -__all__ = [ - "NormType", - "PixelNorm", - "build_normalization_layer", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/normalization.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/normalization.py deleted file mode 100644 index 8fd9adaf9..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/normalization.py +++ /dev/null @@ -1,57 +0,0 @@ -from enum import Enum - -import torch -from torch import nn - - -class NormType(Enum): - """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm).""" - - GROUP = "group" - PIXEL = "pixel" - - -class PixelNorm(nn.Module): - """ - Per-pixel (per-location) RMS normalization layer. - For each element along the chosen dimension, this layer normalizes the tensor - by the root-mean-square of its values across that dimension: - y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps) - """ - - def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: - """ - Args: - dim: Dimension along which to compute the RMS (typically channels). - eps: Small constant added for numerical stability. - """ - super().__init__() - self.dim = dim - self.eps = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Apply RMS normalization along the configured dimension. - """ - # Compute mean of squared values along `dim`, keep dimensions for broadcasting. - mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) - # Normalize by the root-mean-square (RMS). - rms = torch.sqrt(mean_sq + self.eps) - return x / rms - - -def build_normalization_layer(in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP) -> nn.Module: - """ - Create a normalization layer based on the normalization type. - Args: - in_channels: Number of input channels - num_groups: Number of groups for group normalization - normtype: Type of normalization: "group" or "pixel" - Returns: - A normalization layer - """ - if normtype == NormType.GROUP: - return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) - if normtype == NormType.PIXEL: - return PixelNorm(dim=1, eps=1e-6) - raise ValueError(f"Invalid normalization type: {normtype}") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/model_protocol.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/model_protocol.py deleted file mode 100644 index f18e1c3ae..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/model_protocol.py +++ /dev/null @@ -1,39 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Protocol, TypeVar - -import torch - -if TYPE_CHECKING: - from ltx_core.guidance.perturbations import BatchedPerturbationConfig - from ltx_core.model.transformer.modality import Modality - -ModelType = TypeVar("ModelType", covariant=True, bound=torch.nn.Module) # noqa: PLC0105 - - -class ModelConfigurator(Protocol[ModelType]): - """Protocol for model loader classes that instantiates models from a configuration dictionary.""" - - @classmethod - def from_config(cls, config: dict) -> ModelType: ... - - -class LTXModelProtocol(Protocol): - """Velocity-model forward interface shared by ``LTXModel`` and its multi-GPU wrappers. - ``forward`` pins the real signature (enforced structurally); ``__call__`` mirrors it - so protocol-typed values stay callable via ``model(...)``. - """ - - def forward( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ... - - def __call__( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ... diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/__init__.py deleted file mode 100644 index 7d2936b1f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Transformer model components.""" - -from ltx_core.model.transformer.modality import Modality -from ltx_core.model.transformer.model import LTXModel, X0Model -from ltx_core.model.transformer.model_configurator import ( - LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP, - LTXV_MODEL_COMFY_RENAMING_MAP, - LTXAudioOnlyModelConfigurator, - LTXModelConfigurator, - LTXVideoOnlyModelConfigurator, -) - -__all__ = [ - "LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP", - "LTXV_MODEL_COMFY_RENAMING_MAP", - "LTXAudioOnlyModelConfigurator", - "LTXModel", - "LTXModelConfigurator", - "LTXVideoOnlyModelConfigurator", - "Modality", - "X0Model", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/adaln.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/adaln.py deleted file mode 100644 index 65e3622f3..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/adaln.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Optional, Tuple - -import torch -from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings - -# Number of AdaLN modulation parameters per transformer block. -# Base: 2 params (shift + scale) x 3 norms (self-attn, feed-forward, output). -ADALN_NUM_BASE_PARAMS = 6 -# Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm. -ADALN_NUM_CROSS_ATTN_PARAMS = 3 - - -def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: - """Total number of AdaLN parameters per block.""" - return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0) - - -class AdaLayerNormSingle(torch.nn.Module): - r""" - Norm layer adaptive layer norm single (adaLN-single). - As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). - Parameters: - embedding_dim (`int`): The size of each embedding vector. - use_additional_conditions (`bool`): To use additional conditions for normalization or not. - """ - - def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): - super().__init__() - - self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( - embedding_dim, - size_emb_dim=embedding_dim // 3, - ) - - self.silu = torch.nn.SiLU() - self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) - - def forward( - self, - timestep: torch.Tensor, - hidden_dtype: Optional[torch.dtype] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) - return self.linear(self.silu(embedded_timestep)), embedded_timestep diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/attention.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/attention.py deleted file mode 100644 index af14f1131..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/attention.py +++ /dev/null @@ -1,836 +0,0 @@ -import functools -import math -from dataclasses import dataclass, field -from enum import Enum -from typing import Protocol - -import torch -from ltx_core.model.transformer.ops import ( - GatedAttentionCallable, - PreAttentionCallable, - PytorchGatedAttention, - PytorchPreAttention, -) -from ltx_core.model.transformer.rope import LTXRopeType -from torch.nn.attention import SDPBackend, sdpa_kernel -from torch.nn.attention.flex_attention import BlockMask, flex_attention - - -def _torch_default_sdpa_priority() -> list[SDPBackend]: - """Fetch torch's current default SDPA priority order at runtime. - Used as the default for ``PytorchAttention`` so the wrapper-always - code path matches torch's native dispatch order without hard-coding it - (which would drift if torch updates the default). - ``torch._C._get_sdp_priority_order`` is a private API; we accept that - risk because the project pins ``torch`` in the lockfile, so any - rename/removal surfaces on a controlled torch bump rather than silently. - """ - return [SDPBackend(p) for p in torch._C._get_sdp_priority_order()] - - -memory_efficient_attention = None -flash_attn_interface = None -flash_attn_4_func = None -flex_attention_autotuned = flex_attention -flex_attention_default = flex_attention -try: - flex_attention_autotuned = torch.compile(flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") -except Exception: - pass -try: - # LTX cross-attention has large rectangular audio/video sequence shapes. - # Avoid per-rank kernel benchmarking for this path while FSDP collectives - # are queued; square self-attention keeps the same autotuned path as WAN. - flex_attention_default = torch.compile(flex_attention, dynamic=False, mode="default") -except Exception: - pass -try: - from xformers.ops import memory_efficient_attention -except ImportError: - memory_efficient_attention = None -try: - # FlashAttention3 and XFormersAttention cannot be used together - if memory_efficient_attention is None: - import flash_attn_interface -except ImportError: - flash_attn_interface = None -try: - from flash_attn.cute import flash_attn_func as flash_attn_4_func -except (ImportError, AttributeError): - flash_attn_4_func = None - - -class AttentionCallable(Protocol): - """Unmasked attention. Backends without a mask kernel (FA3/FA4) implement only - this protocol; backends that support masks too (Pytorch/SDPA, xFormers) are - structurally usable here and as :class:`MaskedAttentionCallable`.""" - - def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ... - - -class MaskedAttentionCallable(Protocol): - """Masked attention. Mask is required (not optional) -- the caller has already - decided this is the masked path and chosen a backend that can serve it. Used - by :class:`Attention` when its forward receives a non-None ``mask``.""" - - def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor) -> torch.Tensor: ... - - -class PytorchAttention(AttentionCallable): - def __init__(self, priority: list[SDPBackend] | None = None) -> None: - # priority=None -> snapshot torch's default SDPA priority at construction. - # Always passed through ``sdpa_kernel(..., set_priority=True)`` so the - # call site is uniform regardless of how the priority was chosen. - self._priority = priority if priority is not None else _torch_default_sdpa_priority() - - @property - def label(self) -> str: - """Human-readable identifier for this backend. Encodes the SDPA priority - list so a single-backend pin reads differently from the full-priority - dispatcher walk.""" - return f"SDPA[{'>'.join(b.name for b in self._priority)}]" - - def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None) -> torch.Tensor: - b, _, dim_head = q.shape - dim_head //= heads - q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v)) - - if mask is not None: - # add a batch dimension if there isn't already one - if mask.ndim == 2: - mask = mask.unsqueeze(0) - # add a heads dimension if there isn't already one - if mask.ndim == 3: - mask = mask.unsqueeze(1) - - with sdpa_kernel(self._priority, set_priority=True): - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) - out = out.transpose(1, 2).reshape(b, -1, heads * dim_head) - return out - - -class XFormersAttention(AttentionCallable): - label = "xFormers" - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - heads: int, - mask: torch.Tensor | None = None, - ) -> torch.Tensor: - if memory_efficient_attention is None: - raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.") - - b, _, dim_head = q.shape - dim_head //= heads - - # xformers expects [B, M, H, K] - q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) - - if mask is not None: - # add a singleton batch dimension - if mask.ndim == 2: - mask = mask.unsqueeze(0) - # add a singleton heads dimension - if mask.ndim == 3: - mask = mask.unsqueeze(1) - # pad to a multiple of 8 - pad = 8 - mask.shape[-1] % 8 - # the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk) - # but when using separated heads, the shape has to be (B, H, Nq, Nk) - # in flux, this matrix ends up being over 1GB - # here, we create a mask with the same batch/head size as the input mask (potentially singleton or full) - mask_out = torch.empty([mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device) - - mask_out[..., : mask.shape[-1]] = mask - # doesn't this remove the padding again?? - mask = mask_out[..., : mask.shape[-1]] - mask = mask.expand(b, heads, -1, -1) - - out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0) - out = out.reshape(b, -1, heads * dim_head) - return out - - -class FlashAttention3(AttentionCallable): - label = "FlashAttention3" - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - heads: int, - ) -> torch.Tensor: - if flash_attn_interface is None: - raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.") - - b, _, dim_head = q.shape - dim_head //= heads - - q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) - - out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v) - out = out.reshape(b, -1, heads * dim_head) - return out - - -class FlashAttention4(AttentionCallable): - label = "FlashAttention4" - - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - heads: int, - ) -> torch.Tensor: - if flash_attn_4_func is None: - raise RuntimeError("FlashAttention4 was selected but `flash-attn-4` is not installed.") - - b, _, dim_head = q.shape - dim_head //= heads - - q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v)) - - out, _ = flash_attn_4_func(q.to(v.dtype), k.to(v.dtype), v) - out = out.reshape(b, -1, heads * dim_head) - return out - - -# --- Automatic selection ----------------------------------------------------- -# AUTOMATIC inspects installed extras and the GPU arch and returns the fastest -# usable callable for each path. The selection runs once per process (cached). -# The unmasked and masked picks are independent: each calls its own helper and -# may end up on different backends (e.g. FA3 unmasked + xFormers masked on H100). - - -def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool: - """Ask torch whether *backend* can run with the given mask shape. - ``MATH`` is the universal SDPA fallback (pure PyTorch ops, no kernel - requirements) so it returns True everywhere, CPU included. The other - backends use ``torch.backends.cuda.can_use_*`` capability checks (no GPU - compute, no synchronization) and are False without CUDA. The probe shapes - are small but realistic enough to surface constraints (head dim, dtype) - that the per-backend rules care about. - """ - if backend is SDPBackend.MATH: - return True - if not torch.cuda.is_available(): - return False - q = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16) - k = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16) - v = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16) - mask = torch.zeros(1, 4, 128, 128, device="cuda", dtype=torch.bfloat16) if with_mask else None - params = torch.backends.cuda.SDPAParams(q, k, v, mask, 0.0, False, False) - if backend is SDPBackend.CUDNN_ATTENTION: - return torch.backends.cuda.can_use_cudnn_attention(params, debug=False) - if backend is SDPBackend.FLASH_ATTENTION: - return torch.backends.cuda.can_use_flash_attention(params, debug=False) - if backend is SDPBackend.EFFICIENT_ATTENTION: - return torch.backends.cuda.can_use_efficient_attention(params, debug=False) - return False - - -_SDPA_FULL_PRIORITY: tuple[SDPBackend, ...] = ( - SDPBackend.CUDNN_ATTENTION, - SDPBackend.FLASH_ATTENTION, - SDPBackend.EFFICIENT_ATTENTION, - SDPBackend.MATH, -) - - -def _sdpa_full_priority() -> PytorchAttention: - """Hand SDPA the full backend priority order; let torch's dispatcher pick at call time. - ``sdpa_kernel(_SDPA_FULL_PRIORITY, set_priority=True)`` enables all four - backends and orders them; torch then walks the order at call time and picks - the first backend whose ``can_use_*`` check passes for the actual - shapes/dtype/mask. FLASH is rejected automatically when a mask is present; - CUDNN may be rejected under deterministic mode; MATH is the universal - fallback. Probing per-backend usability up front from generic probe shapes - cannot anticipate the variety of real call sites (e.g. broadcast key-only - masks, large head dim), so we defer the choice to the dispatcher. - """ - return PytorchAttention(priority=list(_SDPA_FULL_PRIORITY)) - - -def _select_primary_attention() -> AttentionCallable: - """Pick the fastest unmasked attention based on installed extras and GPU arch. - Priority by arch: - - Hopper (sm_90, H100): FA3 / xFormers (mutually exclusive at import) > FA4 > SDPA. - - Datacenter Blackwell (sm_100, B200): FA4 > SDPA. FA4 is intentionally *not* - picked on consumer Blackwell (sm_120) -- known regressions in newer - FA4 betas; users who want it on sm_120 must opt in explicitly. - - Everywhere else (Ada, Ampere, CPU): SDPA with the full backend priority - list -- torch's runtime dispatcher picks the best fit at call time. - """ - if torch.cuda.is_available(): - major, _ = torch.cuda.get_device_capability(0) - if major == 9: - if flash_attn_interface is not None: - return FlashAttention3() - if memory_efficient_attention is not None: - return XFormersAttention() - if flash_attn_4_func is not None: - return FlashAttention4() - if major == 10 and flash_attn_4_func is not None: - return FlashAttention4() - return _sdpa_full_priority() - - -def _select_masked_attention() -> MaskedAttentionCallable: - """Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with - the full priority list (the dispatcher rejects FLASH automatically when a - mask is present and walks past it).""" - if memory_efficient_attention is not None: - return XFormersAttention() - return _sdpa_full_priority() - - -@functools.cache -def automatic_attention() -> AttentionCallable: - """Cached AUTOMATIC pick for the unmasked path. - Cached so every ``AttentionOps`` in the process shares one instance.""" - return _select_primary_attention() - - -@functools.cache -def automatic_masked_attention() -> MaskedAttentionCallable: - """Cached AUTOMATIC pick for the masked path. See :func:`automatic_attention`.""" - return _select_masked_attention() - - -def attention_label(fn: AttentionCallable | MaskedAttentionCallable) -> str: - """Best-effort human-readable backend name. - Built-in callables expose ``.label`` (encoding the SDPA priority list for the - Pytorch backends); fall back to the class name for custom or wrapped callables - (e.g. the multi-GPU All2All wrappers) that don't define one.""" - return getattr(fn, "label", type(fn).__name__) - - -def _heads_view(x: torch.Tensor, heads: int) -> torch.Tensor: - b, t, inner_dim = x.shape - return x.view(b, t, heads, inner_dim // heads) - - -def _flatten_heads(x: torch.Tensor) -> torch.Tensor: - return x.flatten(2) - - -def _sdpa_shaped(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - out = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2), - k.transpose(1, 2), - v.transpose(1, 2), - attn_mask=None, - dropout_p=0.0, - is_causal=False, - ) - return out.transpose(1, 2) - - -def _flex_attention_flat( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - heads: int, - block_mask: BlockMask, - *, - use_autotune: bool, -) -> torch.Tensor: - q = _heads_view(q, heads) - k = _heads_view(k, heads) - v = _heads_view(v, heads) - q_padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] - kv_padded_length = math.ceil(k.shape[1] / 128) * 128 - k.shape[1] - if q_padded_length > 0: - pad_shape = (q.shape[0], q_padded_length, q.shape[2], q.shape[3]) - q = torch.cat([q, torch.zeros(pad_shape, device=q.device, dtype=q.dtype)], dim=1) - if kv_padded_length > 0: - pad_shape = (k.shape[0], kv_padded_length, k.shape[2], k.shape[3]) - k = torch.cat([k, torch.zeros(pad_shape, device=k.device, dtype=k.dtype)], dim=1) - v = torch.cat([v, torch.zeros(pad_shape, device=v.device, dtype=v.dtype)], dim=1) - flex_fn = flex_attention_autotuned if use_autotune else flex_attention_default - out = flex_fn(query=q.transpose(1, 2), key=k.transpose(1, 2), value=v.transpose(1, 2), block_mask=block_mask) - if q_padded_length > 0: - out = out[:, :, :-q_padded_length] - return out.transpose(1, 2).flatten(2) - - -def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention: - """Build a single-backend ``PytorchAttention`` pin, raising if the backend - can't actually serve the call on this machine. Used by both - :meth:`AttentionFunction.to_callable` and :meth:`MaskedAttentionFunction.to_callable`; - ``with_mask`` differs between the two so the capability check considers - the protocol the caller intends to use. Not used for ``MATH`` -- MATH is - the universal fallback and would falsely fail the CUDA-only probe on CPU. - """ - if not _sdpa_can_use(backend, with_mask=with_mask): - raise RuntimeError( - f"{name} selected but the SDPA {backend.name} backend is not usable on this machine " - "(either no CUDA, the backend rejected the probe shapes, or " - "torch.use_deterministic_algorithms(True) excluded it)." - ) - return PytorchAttention(priority=[backend]) - - -class AttentionFunction(Enum): - PYTORCH = "pytorch" - XFORMERS = "xformers" - FLASH_ATTENTION_3 = "flash_attention_3" - FLASH_ATTENTION_4 = "flash_attention_4" - SDPA_CUDNN = "sdpa_cudnn" - SDPA_FLASH = "sdpa_flash" - SDPA_EFFICIENT = "sdpa_efficient" - SDPA_MATH = "sdpa_math" - # Pick the fastest unmasked backend for the current GPU/extras combo; see - # :func:`automatic_attention`. Default for :class:`AttentionOps`. - AUTOMATIC = "automatic" - - def to_callable(self) -> AttentionCallable: # noqa: PLR0911 - """Resolve to a concrete callable. Use this at module init time so that - torch.compile can trace through the attention call without graph breaks. - Every non-AUTOMATIC variant raises :class:`RuntimeError` when the backend - isn't usable on this machine -- missing package or SDPA backend rejected - on this hardware (e.g. cuDNN under ``torch.use_deterministic_algorithms``). - Opting in means "this kernel or fail loudly". ``AUTOMATIC`` returns the - cached :func:`automatic_attention` instance so every build shares one callable. - """ - match self: - case AttentionFunction.AUTOMATIC: - return automatic_attention() - case AttentionFunction.PYTORCH: - return PytorchAttention() - case AttentionFunction.XFORMERS: - if memory_efficient_attention is None: - raise RuntimeError("AttentionFunction.XFORMERS selected but `xformers` is not installed.") - return XFormersAttention() - case AttentionFunction.FLASH_ATTENTION_3: - if flash_attn_interface is None: - raise RuntimeError("AttentionFunction.FLASH_ATTENTION_3 selected but `flash-attn-3` is not installed.") - return FlashAttention3() - case AttentionFunction.FLASH_ATTENTION_4: - if flash_attn_4_func is None: - raise RuntimeError("AttentionFunction.FLASH_ATTENTION_4 selected but `flash-attn-4` is not installed.") - return FlashAttention4() - case AttentionFunction.SDPA_MATH: - return PytorchAttention(priority=[SDPBackend.MATH]) - case AttentionFunction.SDPA_CUDNN: - return _resolve_sdpa_variant(SDPBackend.CUDNN_ATTENTION, "AttentionFunction.SDPA_CUDNN", with_mask=False) - case AttentionFunction.SDPA_FLASH: - return _resolve_sdpa_variant(SDPBackend.FLASH_ATTENTION, "AttentionFunction.SDPA_FLASH", with_mask=False) - case AttentionFunction.SDPA_EFFICIENT: - return _resolve_sdpa_variant(SDPBackend.EFFICIENT_ATTENTION, "AttentionFunction.SDPA_EFFICIENT", with_mask=False) - - -class MaskedAttentionFunction(Enum): - """Backends usable on the masked path. Mirrors :class:`AttentionFunction` minus - the variants the torch SDPA dispatcher (or the wrapped kernel) rejects with a - mask: ``SDPA_FLASH`` -- FLASH kernel cannot serve an additive ``attn_mask``; - ``FLASH_ATTENTION_3``/``FLASH_ATTENTION_4`` -- neither has a mask kernel at all. - Keeping them out makes "this backend cannot mask" a type error, not a runtime one.""" - - PYTORCH = "pytorch" - XFORMERS = "xformers" - SDPA_CUDNN = "sdpa_cudnn" - SDPA_EFFICIENT = "sdpa_efficient" - SDPA_MATH = "sdpa_math" - # Pick the fastest mask-capable backend for the current extras combo; see - # :func:`automatic_masked_attention`. Default for the masked slot of - # :class:`AttentionOps`. - AUTOMATIC = "automatic" - - def to_callable(self) -> MaskedAttentionCallable: - """Resolve to a concrete masked callable. Same backend classes as - :meth:`AttentionFunction.to_callable`; the protocol returned just exposes - the masked call signature. - Non-AUTOMATIC variants raise :class:`RuntimeError` when the backend isn't - usable for the masked path on this machine. SDPA probes run with - ``with_mask=True`` so the capability check considers the protocol the - caller will actually use.""" - match self: - case MaskedAttentionFunction.AUTOMATIC: - return automatic_masked_attention() - case MaskedAttentionFunction.PYTORCH: - return PytorchAttention() - case MaskedAttentionFunction.XFORMERS: - if memory_efficient_attention is None: - raise RuntimeError("MaskedAttentionFunction.XFORMERS selected but `xformers` is not installed.") - return XFormersAttention() - case MaskedAttentionFunction.SDPA_MATH: - return PytorchAttention(priority=[SDPBackend.MATH]) - case MaskedAttentionFunction.SDPA_CUDNN: - return _resolve_sdpa_variant(SDPBackend.CUDNN_ATTENTION, "MaskedAttentionFunction.SDPA_CUDNN", with_mask=True) - case MaskedAttentionFunction.SDPA_EFFICIENT: - return _resolve_sdpa_variant(SDPBackend.EFFICIENT_ATTENTION, "MaskedAttentionFunction.SDPA_EFFICIENT", with_mask=True) - - -@dataclass(frozen=True) -class AttentionOps: - """Pluggable callables consumed by :class:`Attention`.""" - - attention_function: AttentionCallable = field(default_factory=lambda: AttentionFunction.AUTOMATIC.to_callable()) - masked_attention_function: MaskedAttentionCallable = field(default_factory=lambda: MaskedAttentionFunction.AUTOMATIC.to_callable()) - preattention_function: PreAttentionCallable = field(default_factory=PytorchPreAttention) - gated_attention_function: GatedAttentionCallable = field(default_factory=PytorchGatedAttention) - - -class Attention(torch.nn.Module): - def __init__( - self, - query_dim: int, - context_dim: int | None = None, - heads: int = 8, - dim_head: int = 64, - norm_eps: float = 1e-6, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - ops: AttentionOps | None = None, - apply_gated_attention: bool = False, - ) -> None: - super().__init__() - if ops is None: - ops = AttentionOps() - self.rope_type = rope_type - self.attention_function = ops.attention_function - self.masked_attention_function = ops.masked_attention_function - self.preattention_function = ops.preattention_function - self.gated_attention_function = ops.gated_attention_function - - inner_dim = dim_head * heads - context_dim = query_dim if context_dim is None else context_dim - - self.heads = heads - self.dim_head = dim_head - - self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) - - self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) - self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) - self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) - - # Optional per-head gating - if apply_gated_attention: - self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) - else: - self.to_gate_logits = None - - self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) - - def forward( - self, - x: torch.Tensor, - context: torch.Tensor | None = None, - mask: torch.Tensor | BlockMask | None = None, - pe: torch.Tensor | None = None, - k_pe: torch.Tensor | None = None, - perturbation_mask: torch.Tensor | None = None, - all_perturbed: bool = False, - kv_cache: dict | None = None, - current_start: int = 0, - cache_start: int | None = None, - defer_cache_update: bool = False, - detach_cache_update: bool = False, - ) -> torch.Tensor: - """Multi-head attention with optional RoPE, perturbation masking, and per-head gating. - When ``perturbation_mask`` is all zeros, the expensive query/key path - (linear projections, RMSNorm, RoPE) is skipped entirely and only the - value projection is used as a pass-through. - Args: - x: Query input tensor of shape ``(B, T, query_dim)``. - context: Key/value context tensor of shape ``(B, S, context_dim)``. - Falls back to ``x`` (self-attention) when *None*. - mask: Optional attention mask. Interpretation depends on the attention - backend (additive bias for xformers/PyTorch SDPA). A non-None - ``mask`` routes to ``masked_attention_function``; ``None`` keeps - the unmasked path. - pe: Rotary positional embeddings applied to both ``q`` and ``k``. - k_pe: Separate rotary positional embeddings for ``k`` only. When - *None*, ``pe`` is reused for keys. - perturbation_mask: Optional mask in ``[0, 1]`` that - blends the attention output with the raw value projection: - ``out = attn_out * mask + v * (1 - mask)``. - **1** keeps the full attention output, **0** bypasses attention - and passes the value projection through unchanged. - *None* or all-ones means standard attention; all-zeros skips - the query/key path entirely for efficiency. - all_perturbed: Whether all perturbations are active for this block. - Returns: - Output tensor of shape ``(B, T, query_dim)``. - """ - is_self_attention = context is None - context = x if context is None else context - use_attention = not all_perturbed - - v = self.to_v(context) - - if not use_attention: - out = v - else: - if kv_cache is not None and not is_self_attention: - raise ValueError("LTX KV cache is only supported for self-attention.") - if kv_cache is not None and mask is not None: - raise ValueError("LTX KV cache and explicit self-attention masks are mutually exclusive.") - if cache_start is None: - cache_start = current_start - q = self.to_q(x) - k = self.to_k(context) - q, k = self.preattention_function(q, k, self, mask, pe, k_pe) - cache_update_info = None - if kv_cache is not None: - out = self._cached_attention( - q, - k, - v, - kv_cache=kv_cache, - current_start=int(current_start), - defer_cache_update=defer_cache_update, - detach_cache_update=detach_cache_update, - ) - if isinstance(out, tuple): - out, cache_update_info = out - elif isinstance(mask, BlockMask): - out = _flex_attention_flat(q, k, v, self.heads, mask, use_autotune=is_self_attention) - elif mask is None: - out = self.attention_function(q, k, v, self.heads) # (B, T, H*D) - else: - out = self.masked_attention_function(q, k, v, self.heads, mask) - - if perturbation_mask is not None: - out = out * perturbation_mask + v * (1 - perturbation_mask) - - # Apply per-head gating if enabled - if self.to_gate_logits is not None: - out = self.gated_attention_function(x, out, self) - - out = self.to_out(out) - if use_attention and kv_cache is not None and cache_update_info is not None: - return out, cache_update_info - return out - - def _cached_attention( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - kv_cache: dict, - current_start: int, - defer_cache_update: bool, - detach_cache_update: bool, - ) -> torch.Tensor | tuple[torch.Tensor, tuple[int, int, dict | None]]: - q = _heads_view(q, self.heads) - k = _heads_view(k, self.heads) - v = _heads_view(v, self.heads) - if kv_cache["k"].shape[2] != k.shape[2]: - raise ValueError(f"LTX KV cache stores {kv_cache['k'].shape[2]} heads, but current key has {k.shape[2]} heads.") - - current_end = current_start + q.shape[1] - kv_cache_size = int(kv_cache["k"].shape[1]) - attention_window_size = int(kv_cache.get("attention_window_size", kv_cache_size)) - sink_tokens = int(kv_cache.get("sink_tokens", 0)) - num_new_tokens = int(q.shape[1]) - cache_update_info = None - detach_cache_update = bool(detach_cache_update and torch.is_grad_enabled() and (k.requires_grad or v.requires_grad) and not defer_cache_update) - - local_end_index = int(kv_cache["local_end_index"].item()) + current_end - int(kv_cache["global_end_index"].item()) - local_start_index = local_end_index - num_new_tokens - needs_roll = current_end > int(kv_cache["global_end_index"].item()) and num_new_tokens + int(kv_cache["local_end_index"].item()) > kv_cache_size - - if needs_roll: - num_evicted_tokens = num_new_tokens + int(kv_cache["local_end_index"].item()) - kv_cache_size - num_rolled_tokens = int(kv_cache["local_end_index"].item()) - num_evicted_tokens - sink_tokens - local_end_index -= num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - window_start = max(0, local_end_index - attention_window_size) - if defer_cache_update: - window_k, window_v = self._deferred_window_after_roll( - kv_cache, - k, - v, - window_start, - sink_tokens, - num_rolled_tokens, - num_evicted_tokens, - local_start_index, - num_new_tokens, - ) - cache_update_info = { - "action": "roll_and_insert", - "sink_tokens": sink_tokens, - "num_rolled_tokens": num_rolled_tokens, - "num_evicted_tokens": num_evicted_tokens, - "local_start_index": local_start_index, - "local_end_index": local_end_index, - "new_k": k, - "new_v": v, - } - elif detach_cache_update: - with torch.no_grad(): - if num_rolled_tokens > 0: - kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - kv_cache["k"][:, local_start_index:local_end_index] = k.detach() - kv_cache["v"][:, local_start_index:local_end_index] = v.detach() - window_k, window_v = self._window_with_current_kv( - kv_cache, - k, - v, - window_start, - local_start_index, - num_new_tokens, - ) - else: - self._roll_cache( - kv_cache, - sink_tokens, - num_rolled_tokens, - num_evicted_tokens, - local_start_index, - local_end_index, - k, - v, - no_grad=False, - ) - window_k = kv_cache["k"][:, window_start:local_end_index] - window_v = kv_cache["v"][:, window_start:local_end_index] - else: - window_start = max(0, local_end_index - attention_window_size) - if defer_cache_update: - window_k, window_v = self._deferred_window( - kv_cache, - k, - v, - window_start, - local_start_index, - num_new_tokens, - ) - cache_update_info = { - "action": "direct_insert", - "local_start_index": local_start_index, - "local_end_index": local_end_index, - "new_k": k, - "new_v": v, - } - else: - if detach_cache_update: - with torch.no_grad(): - kv_cache["k"][:, local_start_index:local_end_index] = k.detach() - kv_cache["v"][:, local_start_index:local_end_index] = v.detach() - window_k, window_v = self._window_with_current_kv( - kv_cache, - k, - v, - window_start, - local_start_index, - num_new_tokens, - ) - else: - kv_cache["k"][:, local_start_index:local_end_index] = k - kv_cache["v"][:, local_start_index:local_end_index] = v - window_k = kv_cache["k"][:, window_start:local_end_index] - window_v = kv_cache["v"][:, window_start:local_end_index] - - out = _flatten_heads(_sdpa_shaped(q, window_k, window_v)) - if not defer_cache_update: - kv_cache["global_end_index"].fill_(current_end) - kv_cache["local_end_index"].fill_(local_end_index) - return out - return out, (current_end, local_end_index, cache_update_info) - - @staticmethod - def _window_with_current_kv(kv_cache, current_k, current_v, window_start, local_start_index, num_new_tokens): - window_k_pieces = [] - window_v_pieces = [] - if window_start < local_start_index: - window_k_pieces.append(kv_cache["k"][:, window_start:local_start_index]) - window_v_pieces.append(kv_cache["v"][:, window_start:local_start_index]) - current_start = max(window_start, local_start_index) - local_start_index - if current_start < num_new_tokens: - window_k_pieces.append(current_k[:, current_start:]) - window_v_pieces.append(current_v[:, current_start:]) - return ( - torch.cat(window_k_pieces, dim=1) if len(window_k_pieces) > 1 else window_k_pieces[0], - torch.cat(window_v_pieces, dim=1) if len(window_v_pieces) > 1 else window_v_pieces[0], - ) - - @staticmethod - def _deferred_window(kv_cache, new_k, new_v, window_start, local_start_index, num_new_tokens): - window_k_pieces = [] - window_v_pieces = [] - if window_start < local_start_index: - window_k_pieces.append(kv_cache["k"][:, window_start:local_start_index]) - window_v_pieces.append(kv_cache["v"][:, window_start:local_start_index]) - new_start = max(window_start, local_start_index) - local_start_index - if new_start < num_new_tokens: - window_k_pieces.append(new_k[:, new_start:]) - window_v_pieces.append(new_v[:, new_start:]) - return ( - torch.cat(window_k_pieces, dim=1) if len(window_k_pieces) > 1 else window_k_pieces[0], - torch.cat(window_v_pieces, dim=1) if len(window_v_pieces) > 1 else window_v_pieces[0], - ) - - @staticmethod - def _deferred_window_after_roll( - kv_cache, - new_k, - new_v, - window_start, - sink_tokens, - num_rolled_tokens, - num_evicted_tokens, - local_start_index, - num_new_tokens, - ): - window_k_pieces = [] - window_v_pieces = [] - prefix_end = min(sink_tokens, local_start_index) - if window_start < prefix_end: - window_k_pieces.append(kv_cache["k"][:, window_start:prefix_end]) - window_v_pieces.append(kv_cache["v"][:, window_start:prefix_end]) - rolled_start = max(window_start, sink_tokens) - if rolled_start < local_start_index and num_rolled_tokens > 0: - src_start = rolled_start + num_evicted_tokens - src_end = local_start_index + num_evicted_tokens - window_k_pieces.append(kv_cache["k"][:, src_start:src_end]) - window_v_pieces.append(kv_cache["v"][:, src_start:src_end]) - new_start = max(window_start, local_start_index) - local_start_index - if new_start < num_new_tokens: - window_k_pieces.append(new_k[:, new_start:]) - window_v_pieces.append(new_v[:, new_start:]) - return ( - torch.cat(window_k_pieces, dim=1) if len(window_k_pieces) > 1 else window_k_pieces[0], - torch.cat(window_v_pieces, dim=1) if len(window_v_pieces) > 1 else window_v_pieces[0], - ) - - @staticmethod - def _roll_cache( - kv_cache, - sink_tokens, - num_rolled_tokens, - num_evicted_tokens, - local_start_index, - local_end_index, - new_k, - new_v, - *, - no_grad, - ): - def _assign(): - if num_rolled_tokens > 0: - kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][:, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][:, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["k"][:, local_start_index:local_end_index] = new_k - kv_cache["v"][:, local_start_index:local_end_index] = new_v - - if no_grad: - with torch.no_grad(): - _assign() - else: - _assign() diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/compiling.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/compiling.py deleted file mode 100644 index 8064d3d09..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/compiling.py +++ /dev/null @@ -1,126 +0,0 @@ -from dataclasses import dataclass, field -from typing import Any - -import torch -from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.transformer.model import LTXModel -from ltx_core.model.transformer.transformer_args import BlockPerturbationsProcessor, TransformerArgs - -# Defaults applied inside the patched forward. Overriding via CompilationConfig -# replaces these wholesale; it does not merge. -_DEFAULT_INDUCTOR_CONFIG: dict[str, Any] = {} -_DEFAULT_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256} - - -@dataclass(frozen=True) -class CompilationConfig: - """``torch.compile`` configuration for transformer blocks. ``None`` keeps eager.""" - - mode: str | None = None - backend: str = "inductor" - fullgraph: bool = False - dynamic: bool | None = None - inductor_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_INDUCTOR_CONFIG)) - dynamo_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_DYNAMO_CONFIG)) - - -class _SeqDynamicMarkingProcessor: - """Marks the per-block seq dim dynamic, then delegates to an inner processor. - Installed by ``compile_transformer`` so the per-block compile artifact stays - shape-polymorphic. Wraps whatever ``block_input_processor`` was already on - the model -- callers that customised the processor keep their customisation; - only the seq-dim marking is layered on top. Lives outside the compiled - region, so ``mark_dynamic`` runs in eager mode on the tensors that are - about to cross into the trace. - """ - - def __init__(self, inner: BlockPerturbationsProcessor) -> None: - self.inner = inner - - def __call__( - self, - args: TransformerArgs, - perturbations: BatchedPerturbationConfig, - block_idx: int, - self_attn_type: PerturbationType, - cross_attn_type: PerturbationType, - ) -> TransformerArgs: - # Positional embeddings are second-from-last regardless of rope type: - # split rope is (B, H, T, D//2) -- dim -2 == 2; interleaved rope is (B, T, D) - # -- dim -2 == 1. Both work via the negative index. - torch._dynamo.mark_dynamic(args.x, 1) - cos, sin = args.positional_embeddings - torch._dynamo.mark_dynamic(cos, cos.ndim - 2) - torch._dynamo.mark_dynamic(sin, sin.ndim - 2) - if args.cross_positional_embeddings is not None: - cross_cos, cross_sin = args.cross_positional_embeddings - torch._dynamo.mark_dynamic(cross_cos, cross_cos.ndim - 2) - torch._dynamo.mark_dynamic(cross_sin, cross_sin.ndim - 2) - if args.self_attention_mask is not None: - # Dense form is (B, 1, T, T); key-padding form (from the SP wrapper) - # is (B, 1, 1, T) -- leave the size-1 query dim static so Dynamo - # keeps the broadcast. - if args.self_attention_mask.shape[2] > 1: - torch._dynamo.mark_dynamic(args.self_attention_mask, 2) - torch._dynamo.mark_dynamic(args.self_attention_mask, 3) - if args.context_mask is not None: - torch._dynamo.mark_dynamic(args.context_mask, 2) - # `timesteps` / `embedded_timestep` are per-token when conditioning sets a - # per-position denoise mask, in which case their dim 1 equals the seq length - # and must vary with it. When they're a single timestep broadcast across the - # sequence (dim 1 == 1), leaving them static lets Dynamo keep the size-1 - # broadcast. - if args.timesteps.shape[1] > 1: - torch._dynamo.mark_dynamic(args.timesteps, 1) - if args.embedded_timestep.shape[1] > 1: - torch._dynamo.mark_dynamic(args.embedded_timestep, 1) - # `cross_scale_shift_timestep` is the cross-attn AdaLN scale/shift input - # derived from the own-modality per-token timesteps (denoise_mask * sigma), - # so its dim 1 equals the seq length when conditioning is per-token. - # `cross_gate_timestep` is the cross-modality sigma scalar -- dim 1 is 1 - # and broadcasts, leave it static. Same guard pattern as `timesteps`. - if args.cross_scale_shift_timestep is not None and args.cross_scale_shift_timestep.shape[1] > 1: - torch._dynamo.mark_dynamic(args.cross_scale_shift_timestep, 1) - return self.inner(args, perturbations, block_idx, self_attn_type, cross_attn_type) - - -def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel: - """Compile each transformer block via ``torch.compile`` with the given settings. - The patched forward emits ``torch.compiler.cudagraph_mark_step_begin()`` once - per step. Under CUDA-graph-enabling modes (``"reduce-overhead"`` / - ``"max-autotune"``) this overrides Dynamo's per-invocation auto-mark - heuristic, which would otherwise fire once per compiled block call (48 per - forward) and treat each block call as a fresh iteration. Under other modes - the mark is a no-op (decrements an unread counter). - """ - model.transformer_blocks = torch.nn.ModuleList(torch.compile(m, mode=config.mode, backend=config.backend, fullgraph=config.fullgraph, dynamic=config.dynamic) for m in model.transformer_blocks) - model.block_input_processor = _SeqDynamicMarkingProcessor(inner=model.block_input_processor) - - def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: - torch.compiler.cudagraph_mark_step_begin() - with ( - torch._inductor.config.patch(**config.inductor_config), - torch._dynamo.config.patch(**config.dynamo_config), # type: ignore[attr-defined] - ): - return model.forward_without_compilation(*args, **kwargs) - - model.forward_without_compilation = model.forward - model.forward = patched_dynamo_forward - return model - - -def build_compile_transformer_op(config: CompilationConfig) -> ModuleOps: - """Build a ``ModuleOps`` that compiles transformer blocks with the given settings.""" - return ModuleOps( - name="compile_transformer", - matcher=lambda model: isinstance(model, LTXModel), - mutator=lambda model: compile_transformer(model, config), - ) - - -def modify_sd_ops_for_compilation(original_sd_ops: SDOps, number_of_blocks: int = 48) -> SDOps: - for i in range(number_of_blocks): - original_sd_ops = original_sd_ops.with_replacement(f"transformer_blocks.{i}.", f"transformer_blocks.{i}._orig_mod.") - return original_sd_ops diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/feed_forward.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/feed_forward.py deleted file mode 100644 index 3ee404aeb..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/feed_forward.py +++ /dev/null @@ -1,14 +0,0 @@ -import torch -from ltx_core.model.transformer.gelu_approx import GELUApprox - - -class FeedForward(torch.nn.Module): - def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: - super().__init__() - inner_dim = int(dim * mult) - project_in = GELUApprox(dim, inner_dim) - - self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/gelu_approx.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/gelu_approx.py deleted file mode 100644 index 1923e6e22..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/gelu_approx.py +++ /dev/null @@ -1,10 +0,0 @@ -import torch - - -class GELUApprox(torch.nn.Module): - def __init__(self, dim_in: int, dim_out: int) -> None: - super().__init__() - self.proj = torch.nn.Linear(dim_in, dim_out) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.gelu(self.proj(x), approximate="tanh") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/modality.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/modality.py deleted file mode 100644 index 0e213b142..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/modality.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -import dataclasses -from dataclasses import dataclass - -import torch - - -@dataclass(frozen=True) -class Modality: - """ - Input data for a single modality (video or audio) in the transformer. - Bundles the latent tokens, timestep embeddings, positional information, - and text conditioning context for processing by the diffusion transformer. - Attributes: - latent: Patchified latent tokens, shape ``(B, T, D)`` where *B* is - the batch size, *T* is the total number of tokens (noisy + - conditioning), and *D* is the input dimension. - timesteps: Per-token timestep embeddings, shape ``(B, T)``. - positions: Per-token patch coordinates used to build the RoPE - frequencies. With the default ``use_middle_indices_grid=True``, - shape is ``(B, n_pos_dims, T, 2)`` where ``n_pos_dims=3`` for - video (time, height, width) and ``n_pos_dims=1`` for audio - (time); the last dim of size 2 holds the ``[start, end)`` - index bounds of each patch, and RoPE is evaluated at the - *middle* of that range -- hence the flag name. Taking the - patch midpoint produces a smoother and more accurate - positional signal than indexing by the patch's start when - patches span more than one spatial / temporal unit. - When ``use_middle_indices_grid=False``, the legacy 3-D form - ``(B, n_pos_dims, T)`` of integer positional indices is - accepted instead and used as-is (no midpoint derivation). - context: Text conditioning embeddings from the prompt encoder. - enabled: Whether this modality is active in the current forward pass. - context_mask: Optional mask for the text context tokens. - attention_mask: Optional 2-D self-attention mask, shape ``(B, T, T)``. - Values in ``[0, 1]`` where ``1`` = full attention and ``0`` = no - attention. ``None`` means unrestricted (full) attention between - all tokens. Built incrementally by conditioning items; see - :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`. - """ - - latent: torch.Tensor # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension - sigma: torch.Tensor # Shape: (B,). Current sigma value, used for cross-attention timestep calculation. - timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps - # Shape: (B, n_pos_dims, T, 2) by default (use_middle_indices_grid=True); - # n_pos_dims=3 for video, 1 for audio; last dim holds [start, end) patch bounds. - # Legacy form (B, n_pos_dims, T) when use_middle_indices_grid=False. - positions: torch.Tensor - context: torch.Tensor - enabled: bool = True - context_mask: torch.Tensor | None = None - attention_mask: torch.Tensor | None = None - cross_attention_mask: torch.Tensor | None = None - - def split(self, sizes: list[int]) -> list[Modality]: - """Split along the batch dimension into chunks of the given sizes.""" - n = len(sizes) - split_fields: dict[str, list[torch.Tensor | None] | list[bool]] = {} - for f in dataclasses.fields(self): - value = getattr(self, f.name) - if isinstance(value, torch.Tensor): - split_fields[f.name] = list(value.split(sizes, dim=0)) - elif value is None or isinstance(value, bool) or type(value).__name__ == "BlockMask": - split_fields[f.name] = [value] * n - else: - raise TypeError(f"Cannot split field {f.name!r}: unsupported type {type(value)}") - return [Modality(**{name: parts[i] for name, parts in split_fields.items()}) for i in range(n)] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model.py deleted file mode 100644 index 4e797806d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model.py +++ /dev/null @@ -1,769 +0,0 @@ -import logging -import math -from dataclasses import replace -from enum import Enum - -import torch -from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType -from ltx_core.model.model_protocol import LTXModelProtocol -from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient -from ltx_core.model.transformer.attention import attention_label -from ltx_core.model.transformer.modality import Modality -from ltx_core.model.transformer.rope import LTXRopeType -from ltx_core.model.transformer.transformer import ( - DEFAULT_TRANSFORMER_OPS, - BasicAVTransformerBlock, - TransformerConfig, - TransformerOpsConfig, -) -from ltx_core.model.transformer.transformer_args import ( - BlockPerturbationsProcessor, - MultiModalTransformerArgsPreprocessor, - TransformerArgs, - TransformerArgsPreprocessor, -) -from ltx_core.utils import to_denoised -from torch.nn.attention.flex_attention import BlockMask, create_block_mask - -logger = logging.getLogger(__name__) - - -class LTXModelType(Enum): - AudioVideo = "ltx av model" - VideoOnly = "ltx video only model" - AudioOnly = "ltx audio only model" - - def is_video_enabled(self) -> bool: - return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly) - - def is_audio_enabled(self) -> bool: - return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly) - - -class LTXModel(torch.nn.Module): - """ - LTX model transformer implementation. - This class implements the transformer blocks for the LTX model. - """ - - def __init__( # noqa: PLR0913 - self, - *, - model_type: LTXModelType = LTXModelType.AudioVideo, - num_attention_heads: int = 32, - attention_head_dim: int = 128, - in_channels: int = 128, - out_channels: int = 128, - num_layers: int = 48, - cross_attention_dim: int = 4096, - norm_eps: float = 1e-06, - ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS, - positional_embedding_theta: float = 10000.0, - positional_embedding_max_pos: list[int] | None = None, - timestep_scale_multiplier: int = 1000, - use_middle_indices_grid: bool = True, - audio_num_attention_heads: int = 32, - audio_attention_head_dim: int = 64, - audio_in_channels: int = 128, - audio_out_channels: int = 128, - audio_cross_attention_dim: int = 2048, - audio_positional_embedding_max_pos: list[int] | None = None, - av_ca_timestep_scale_multiplier: int = 1, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - double_precision_rope: bool = False, - apply_gated_attention: bool = False, - caption_projection: torch.nn.Module | None = None, - audio_caption_projection: torch.nn.Module | None = None, - cross_attention_adaln: bool = False, - ): - super().__init__() - # Log the attention backends this transformer is built with. Reading the resolved - # ``label`` off the ops reports whatever was selected -- AUTOMATIC, an explicit pin - # (PYTORCH/XFORMERS/FA3/FA4/SDPA_*), or a directly supplied callable -- so this is the - # single source of truth for which kernel a build uses. Fires once per build. - logger.info( - "Building transformer with attention backends -- self: %s, masked: %s", - attention_label(ops.attention_ops.attention_function), - attention_label(ops.attention_ops.masked_attention_function), - ) - self._enable_gradient_checkpointing = False - self.cross_attention_adaln = cross_attention_adaln - self.use_middle_indices_grid = use_middle_indices_grid - self.rope_type = rope_type - self.double_precision_rope = double_precision_rope - self.timestep_scale_multiplier = timestep_scale_multiplier - self.positional_embedding_theta = positional_embedding_theta - self.model_type = model_type - self.num_frame_per_block = 1 - self.local_attn_size = -1 - self.sink_size = 0 - self.defer_kv_cache_updates = False - self.detach_kv_cache_updates = False - cross_pe_max_pos = None - if model_type.is_video_enabled(): - if positional_embedding_max_pos is None: - positional_embedding_max_pos = [20, 2048, 2048] - self.positional_embedding_max_pos = positional_embedding_max_pos - self.num_attention_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - self._init_video( - in_channels=in_channels, - out_channels=out_channels, - norm_eps=norm_eps, - caption_projection=caption_projection, - ) - - if model_type.is_audio_enabled(): - if audio_positional_embedding_max_pos is None: - audio_positional_embedding_max_pos = [20] - self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos - self.audio_num_attention_heads = audio_num_attention_heads - self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim - self._init_audio( - in_channels=audio_in_channels, - out_channels=audio_out_channels, - norm_eps=norm_eps, - caption_projection=audio_caption_projection, - ) - - if model_type.is_video_enabled() and model_type.is_audio_enabled(): - cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0]) - self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier - self.audio_cross_attention_dim = audio_cross_attention_dim - self._init_audio_video(num_scale_shift_values=4) - - self._init_preprocessors(cross_pe_max_pos) - # Initialize transformer blocks - self._init_transformer_blocks( - num_layers=num_layers, - attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0, - cross_attention_dim=cross_attention_dim, - audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0, - audio_cross_attention_dim=audio_cross_attention_dim, - norm_eps=norm_eps, - ops=ops, - apply_gated_attention=apply_gated_attention, - ) - # Hook for per-block input prep. Compile transforms in `compiling.py` - # wrap (not replace) this with a processor that also marks the seq dim - # dynamic, so any caller customisation here is preserved as the inner. - self.block_input_processor = BlockPerturbationsProcessor() - - @property - def _adaln_embedding_coefficient(self) -> int: - return adaln_embedding_coefficient(self.cross_attention_adaln) - - def _init_video( - self, - in_channels: int, - out_channels: int, - norm_eps: float, - caption_projection: torch.nn.Module | None = None, - ) -> None: - """Initialize video-specific components.""" - # Video input components - self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) - if caption_projection is not None: - self.caption_projection = caption_projection - - self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient) - - self.prompt_adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None - - # Video output components - self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim)) - self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps) - self.proj_out = torch.nn.Linear(self.inner_dim, out_channels) - - def _init_audio( - self, - in_channels: int, - out_channels: int, - norm_eps: float, - caption_projection: torch.nn.Module | None = None, - ) -> None: - """Initialize audio-specific components.""" - - # Audio input components - self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True) - if caption_projection is not None: - self.audio_caption_projection = caption_projection - - self.audio_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=self._adaln_embedding_coefficient, - ) - - self.audio_prompt_adaln_single = AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None - - # Audio output components - self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim)) - self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps) - self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels) - - def _init_audio_video( - self, - num_scale_shift_values: int, - ) -> None: - """Initialize audio-video cross-attention components.""" - self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle( - self.inner_dim, - embedding_coefficient=num_scale_shift_values, - ) - - self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=num_scale_shift_values, - ) - - self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle( - self.inner_dim, - embedding_coefficient=1, - ) - - self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle( - self.audio_inner_dim, - embedding_coefficient=1, - ) - - def _init_preprocessors( - self, - cross_pe_max_pos: int | None = None, - ) -> None: - """Initialize preprocessors for LTX.""" - - if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled(): - self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor( - patchify_proj=self.patchify_proj, - adaln=self.adaln_single, - cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single, - cross_gate_adaln=self.av_ca_a2v_gate_adaln_single, - inner_dim=self.inner_dim, - max_pos=self.positional_embedding_max_pos, - num_attention_heads=self.num_attention_heads, - cross_pe_max_pos=cross_pe_max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - audio_cross_attention_dim=self.audio_cross_attention_dim, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, - caption_projection=getattr(self, "caption_projection", None), - prompt_adaln=getattr(self, "prompt_adaln_single", None), - ) - self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor( - patchify_proj=self.audio_patchify_proj, - adaln=self.audio_adaln_single, - cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single, - cross_gate_adaln=self.av_ca_v2a_gate_adaln_single, - inner_dim=self.audio_inner_dim, - max_pos=self.audio_positional_embedding_max_pos, - num_attention_heads=self.audio_num_attention_heads, - cross_pe_max_pos=cross_pe_max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - audio_cross_attention_dim=self.audio_cross_attention_dim, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier, - caption_projection=getattr(self, "audio_caption_projection", None), - prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), - ) - elif self.model_type.is_video_enabled(): - self.video_args_preprocessor = TransformerArgsPreprocessor( - patchify_proj=self.patchify_proj, - adaln=self.adaln_single, - inner_dim=self.inner_dim, - max_pos=self.positional_embedding_max_pos, - num_attention_heads=self.num_attention_heads, - use_middle_indices_grid=self.use_middle_indices_grid, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - caption_projection=getattr(self, "caption_projection", None), - prompt_adaln=getattr(self, "prompt_adaln_single", None), - ) - elif self.model_type.is_audio_enabled(): - self.audio_args_preprocessor = TransformerArgsPreprocessor( - patchify_proj=self.audio_patchify_proj, - adaln=self.audio_adaln_single, - inner_dim=self.audio_inner_dim, - max_pos=self.audio_positional_embedding_max_pos, - num_attention_heads=self.audio_num_attention_heads, - use_middle_indices_grid=self.use_middle_indices_grid, - timestep_scale_multiplier=self.timestep_scale_multiplier, - double_precision_rope=self.double_precision_rope, - positional_embedding_theta=self.positional_embedding_theta, - rope_type=self.rope_type, - caption_projection=getattr(self, "audio_caption_projection", None), - prompt_adaln=getattr(self, "audio_prompt_adaln_single", None), - ) - - def _init_transformer_blocks( - self, - num_layers: int, - attention_head_dim: int, - cross_attention_dim: int, - audio_attention_head_dim: int, - audio_cross_attention_dim: int, - norm_eps: float, - ops: TransformerOpsConfig, - apply_gated_attention: bool, - ) -> None: - """Initialize transformer blocks for LTX.""" - video_config = ( - TransformerConfig( - dim=self.inner_dim, - heads=self.num_attention_heads, - d_head=attention_head_dim, - context_dim=cross_attention_dim, - apply_gated_attention=apply_gated_attention, - cross_attention_adaln=self.cross_attention_adaln, - ) - if self.model_type.is_video_enabled() - else None - ) - audio_config = ( - TransformerConfig( - dim=self.audio_inner_dim, - heads=self.audio_num_attention_heads, - d_head=audio_attention_head_dim, - context_dim=audio_cross_attention_dim, - apply_gated_attention=apply_gated_attention, - cross_attention_adaln=self.cross_attention_adaln, - ) - if self.model_type.is_audio_enabled() - else None - ) - self.transformer_blocks = torch.nn.ModuleList( - [ - BasicAVTransformerBlock( - video=video_config, - audio=audio_config, - rope_type=self.rope_type, - norm_eps=norm_eps, - ops=ops, - ) - for _ in range(num_layers) - ] - ) - - def set_gradient_checkpointing(self, enable: bool) -> None: - """Enable or disable gradient checkpointing for transformer blocks. - Gradient checkpointing trades compute for memory by recomputing activations - during the backward pass instead of storing them. This can significantly - reduce memory usage at the cost of ~20-30% slower training. - Args: - enable: Whether to enable gradient checkpointing - """ - self._enable_gradient_checkpointing = enable - - def _process_transformer_blocks( - self, - video: TransformerArgs | None, - audio: TransformerArgs | None, - perturbations: BatchedPerturbationConfig | None, - kv_cache: dict | None = None, - video_current_start: int = 0, - audio_current_start: int = 0, - defer_cache_updates: bool | None = None, - detach_cache_updates: bool | None = None, - ) -> tuple[TransformerArgs | None, TransformerArgs | None]: - """Process transformer blocks for LTXAV. - Per-block perturbation masks are precomputed here and attached to each - modality's ``TransformerArgs`` so the block forward has no per-block - identity to specialise on — all blocks share a single Dynamo cache slot. - """ - if perturbations is None: - batch_size = (video or audio).x.shape[0] - perturbations = BatchedPerturbationConfig.empty(batch_size) - if defer_cache_updates is None: - defer_cache_updates = bool(self.defer_kv_cache_updates) - defer_cache_updates = bool(defer_cache_updates and torch.is_grad_enabled()) - if detach_cache_updates is None: - detach_cache_updates = bool(self.detach_kv_cache_updates) - detach_cache_updates = bool(detach_cache_updates and torch.is_grad_enabled() and not defer_cache_updates) - - cache_update_infos = [] - - for block_idx, block in enumerate(self.transformer_blocks): - if video is not None: - video = self.block_input_processor( - video, - perturbations, - block_idx, - self_attn_type=PerturbationType.SKIP_VIDEO_SELF_ATTN, - cross_attn_type=PerturbationType.SKIP_A2V_CROSS_ATTN, - ) - if audio is not None: - audio = self.block_input_processor( - audio, - perturbations, - block_idx, - self_attn_type=PerturbationType.SKIP_AUDIO_SELF_ATTN, - cross_attn_type=PerturbationType.SKIP_V2A_CROSS_ATTN, - ) - - if self._enable_gradient_checkpointing and self.training: - checkpoint_tensors = [] - video_template = None - audio_template = None - if video is not None: - checkpoint_tensors.append(video.x) - video_template = replace(video, x=None) - if audio is not None: - checkpoint_tensors.append(audio.x) - audio_template = replace(audio, x=None) - - video_kv_cache = kv_cache["video"][block_idx] if kv_cache is not None and kv_cache.get("video") is not None else None - audio_kv_cache = kv_cache["audio"][block_idx] if kv_cache is not None and kv_cache.get("audio") is not None else None - - def block_forward( - *tensor_args, - block=block, - video_template=video_template, - audio_template=audio_template, - video_kv_cache=video_kv_cache, - audio_kv_cache=audio_kv_cache, - video_current_start=video_current_start, - audio_current_start=audio_current_start, - defer_cache_updates=defer_cache_updates, - detach_cache_updates=detach_cache_updates, - ): - arg_idx = 0 - video_args = video_template - audio_args = audio_template - if video_args is not None: - video_args = replace(video_args, x=tensor_args[arg_idx]) - arg_idx += 1 - if audio_args is not None: - audio_args = replace(audio_args, x=tensor_args[arg_idx]) - return block( - video=video_args, - audio=audio_args, - video_kv_cache=video_kv_cache, - audio_kv_cache=audio_kv_cache, - video_current_start=video_current_start, - audio_current_start=audio_current_start, - defer_cache_update=defer_cache_updates, - detach_cache_update=detach_cache_updates, - ) - - has_block_mask = any(isinstance(mask, BlockMask) for args in (video, audio) if args is not None for mask in (args.self_attention_mask, args.cross_attention_mask)) - with torch.utils.checkpoint.set_checkpoint_early_stop(not has_block_mask): - block_out = torch.utils.checkpoint.checkpoint( - block_forward, - *checkpoint_tensors, - use_reentrant=False, - ) - else: - block_out = block( - video=video, - audio=audio, - video_kv_cache=kv_cache["video"][block_idx] if kv_cache is not None and kv_cache.get("video") is not None else None, - audio_kv_cache=kv_cache["audio"][block_idx] if kv_cache is not None and kv_cache.get("audio") is not None else None, - video_current_start=video_current_start, - audio_current_start=audio_current_start, - defer_cache_update=defer_cache_updates, - detach_cache_update=detach_cache_updates, - ) - - if isinstance(block_out, tuple) and len(block_out) == 3: - video, audio, cache_update_info = block_out - cache_update_infos.append((block_idx, cache_update_info)) - else: - video, audio = block_out - - if kv_cache is not None and defer_cache_updates and cache_update_infos: - self._apply_cache_updates(kv_cache, cache_update_infos) - return video, audio - - @staticmethod - def _apply_cache_updates(kv_cache, cache_update_infos): - for block_index, update_info in cache_update_infos: - for modality_name in ("video", "audio"): - if update_info.get(modality_name) is None or kv_cache.get(modality_name) is None: - continue - current_end, local_end_index, cache_update_info = update_info[modality_name] - cache = kv_cache[modality_name][block_index] - if cache_update_info is not None: - if cache_update_info["action"] == "roll_and_insert": - sink_tokens = cache_update_info["sink_tokens"] - num_rolled_tokens = cache_update_info["num_rolled_tokens"] - num_evicted_tokens = cache_update_info["num_evicted_tokens"] - if num_rolled_tokens > 0: - cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = cache["k"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = cache["v"][ - :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens - ].clone() - cache["k"][:, cache_update_info["local_start_index"] : cache_update_info["local_end_index"]] = cache_update_info["new_k"] - cache["v"][:, cache_update_info["local_start_index"] : cache_update_info["local_end_index"]] = cache_update_info["new_v"] - cache["global_end_index"].fill_(current_end) - cache["local_end_index"].fill_(local_end_index) - - def _process_output( - self, - scale_shift_table: torch.Tensor, - norm_out: torch.nn.LayerNorm, - proj_out: torch.nn.Linear, - x: torch.Tensor, - embedded_timestep: torch.Tensor, - ) -> torch.Tensor: - """Process output for LTXV.""" - # Apply scale-shift modulation - scale_shift_values = scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] - shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] - - x = norm_out(x) - x = x * (1 + scale) + shift - x = proj_out(x) - return x - - def forward( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - kv_cache: dict | None = None, - current_start: int = 0, - video_current_start: int | None = None, - audio_current_start: int | None = None, - defer_cache_updates: bool | None = None, - detach_cache_updates: bool | None = None, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """ - Forward pass for LTX models. - Returns: - Processed output tensors - """ - if not self.model_type.is_video_enabled() and video is not None: - raise ValueError("Video is not enabled for this model") - if not self.model_type.is_audio_enabled() and audio is not None: - raise ValueError("Audio is not enabled for this model") - - video_args = self.video_args_preprocessor.prepare(video, audio) if video is not None else None - audio_args = self.audio_args_preprocessor.prepare(audio, video) if audio is not None else None - # Process transformer blocks - video_out, audio_out = self._process_transformer_blocks( - video=video_args, - audio=audio_args, - perturbations=perturbations, - kv_cache=kv_cache, - video_current_start=current_start if video_current_start is None else video_current_start, - audio_current_start=current_start if audio_current_start is None else audio_current_start, - defer_cache_updates=defer_cache_updates, - detach_cache_updates=detach_cache_updates, - ) - - # Process output - vx = self._process_output(self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep) if video_out is not None else None - ax = ( - self._process_output( - self.audio_scale_shift_table, - self.audio_norm_out, - self.audio_proj_out, - audio_out.x, - audio_out.embedded_timestep, - ) - if audio_out is not None - else None - ) - return vx, ax - - def _forward_inference( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig | None = None, - kv_cache: dict | None = None, - current_start: int = 0, - video_current_start: int | None = None, - audio_current_start: int | None = None, - defer_cache_updates: bool | None = None, - detach_cache_updates: bool | None = None, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - return self.forward( - video=video, - audio=audio, - perturbations=perturbations, - kv_cache=kv_cache, - current_start=current_start, - video_current_start=video_current_start, - audio_current_start=audio_current_start, - defer_cache_updates=defer_cache_updates, - detach_cache_updates=detach_cache_updates, - ) - - @staticmethod - def prepare_blockwise_causal_mask_from_block_ids( - device: torch.device | str, - block_ids: torch.Tensor, - *, - local_attn_blocks: int = -1, - batch_size: int | None = None, - ) -> BlockMask: - block_ids = block_ids.to(device=device, dtype=torch.long) - total_length = int(block_ids.numel()) - padded_length = math.ceil(total_length / 128) * 128 - total_length - padded_block_ids = torch.cat([block_ids, block_ids.new_zeros(padded_length)], dim=0) - - def attention_mask(b, h, q_idx, kv_idx): - is_real_q = q_idx < total_length - is_real_k = kv_idx < total_length - q_block = padded_block_ids[q_idx] - kv_block = padded_block_ids[kv_idx] - if local_attn_blocks == -1: - visible = kv_block <= q_block - else: - visible = (kv_block <= q_block) & (kv_block >= (q_block - local_attn_blocks + 1)) - return (q_idx == kv_idx) | (is_real_q & is_real_k & visible) - - return create_block_mask( - attention_mask, - B=batch_size, - H=None, - Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, - _compile=True, - device=device, - ) - - @staticmethod - def prepare_teacher_forcing_mask_from_block_ids( - device: torch.device | str, - block_ids: torch.Tensor, - *, - batch_size: int | None = None, - ) -> BlockMask: - block_ids = block_ids.to(device=device, dtype=torch.long) - clean_length = int(block_ids.numel()) - total_length = clean_length * 2 - padded_length = math.ceil(total_length / 128) * 128 - total_length - all_block_ids = torch.cat([block_ids, block_ids, block_ids.new_zeros(padded_length)], dim=0) - - def attention_mask(b, h, q_idx, kv_idx): - is_real_q = q_idx < total_length - is_real_k = kv_idx < total_length - q_is_noisy = q_idx >= clean_length - kv_is_noisy = kv_idx >= clean_length - q_block = all_block_ids[q_idx] - kv_block = all_block_ids[kv_idx] - clean_to_clean = (~q_is_noisy) & (~kv_is_noisy) & (kv_block <= q_block) - noisy_to_clean = q_is_noisy & (~kv_is_noisy) & (kv_block < q_block) - noisy_to_noisy = q_is_noisy & kv_is_noisy & (kv_block == q_block) - return (q_idx == kv_idx) | (is_real_q & is_real_k & (clean_to_clean | noisy_to_clean | noisy_to_noisy)) - - return create_block_mask( - attention_mask, - B=batch_size, - H=None, - Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, - _compile=True, - device=device, - ) - - @staticmethod - def prepare_teacher_forcing_cross_mask_from_block_ids( - device: torch.device | str, - query_block_ids: torch.Tensor, - key_block_ids: torch.Tensor, - *, - batch_size: int | None = None, - ) -> BlockMask: - query_block_ids = query_block_ids.to(device=device, dtype=torch.long) - key_block_ids = key_block_ids.to(device=device, dtype=torch.long) - clean_q_length = int(query_block_ids.numel()) - clean_k_length = int(key_block_ids.numel()) - q_length = clean_q_length * 2 - kv_length = clean_k_length * 2 - q_padded_length = math.ceil(q_length / 128) * 128 - q_length - kv_padded_length = math.ceil(kv_length / 128) * 128 - kv_length - all_query_block_ids = torch.cat([query_block_ids, query_block_ids, query_block_ids.new_zeros(q_padded_length)], dim=0) - all_key_block_ids = torch.cat([key_block_ids, key_block_ids, key_block_ids.new_zeros(kv_padded_length)], dim=0) - - def attention_mask(b, h, q_idx, kv_idx): - is_real_q = q_idx < q_length - is_real_k = kv_idx < kv_length - q_is_noisy = q_idx >= clean_q_length - kv_is_noisy = kv_idx >= clean_k_length - q_block = all_query_block_ids[q_idx] - kv_block = all_key_block_ids[kv_idx] - same_stream = (q_is_noisy & kv_is_noisy) | ((~q_is_noisy) & (~kv_is_noisy)) - same_block = q_block == kv_block - # AR inference keeps independent video/audio self-attention caches; - # A2V and V2A attention only exchange the current chunk. Restricting - # cross attention to the same stream and block makes parallel TF - # expose exactly that cross-modal context. - return is_real_q & is_real_k & same_stream & same_block - - return create_block_mask( - attention_mask, - B=batch_size, - H=None, - Q_LEN=q_length + q_padded_length, - KV_LEN=kv_length + kv_padded_length, - _compile=True, - device=device, - ) - - -class LegacyX0Model(torch.nn.Module): - """ - Legacy X0 model implementation. - Returns fully denoised output based on the velocities produced by the base model. - """ - - def __init__(self, velocity_model: LTXModelProtocol): - super().__init__() - self.velocity_model = velocity_model - - def forward( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - sigma: float, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """ - Denoise the video and audio according to the sigma. - Returns: - Denoised video and audio - """ - vx, ax = self.velocity_model(video, audio, perturbations) - denoised_video = to_denoised(video.latent, vx, sigma) if vx is not None else None - denoised_audio = to_denoised(audio.latent, ax, sigma) if ax is not None else None - return denoised_video, denoised_audio - - -class X0Model(torch.nn.Module): - """ - X0 model implementation. - Returns fully denoised outputs based on the velocities produced by the base model. - Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask. - """ - - def __init__(self, velocity_model: LTXModelProtocol): - super().__init__() - self.velocity_model = velocity_model - - def forward( - self, - video: Modality | None, - audio: Modality | None, - perturbations: BatchedPerturbationConfig, - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """ - Denoise the video and audio according to the sigma. - Returns: - Denoised video and audio - """ - vx, ax = self.velocity_model(video, audio, perturbations) - denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None - denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None - return denoised_video, denoised_audio diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model_configurator.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model_configurator.py deleted file mode 100644 index b52c1812c..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model_configurator.py +++ /dev/null @@ -1,213 +0,0 @@ -import torch -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.model.transformer.model import LTXModel, LTXModelType -from ltx_core.model.transformer.rope import LTXRopeType -from ltx_core.model.transformer.text_projection import create_caption_projection -from ltx_core.model.transformer.transformer import DEFAULT_TRANSFORMER_OPS, TransformerOpsConfig -from ltx_core.utils import check_config_value - - -class LTXModelConfigurator(ModelConfigurator[LTXModel]): - """ - Configurator for LTX model. - Used to create an LTX model from a configuration dictionary. - """ - - @classmethod - def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel: - # Build caption projections for 19B models (projection handled in transformer). - caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True) - - config = config.get("transformer", {}) - - check_config_value(config, "dropout", 0.0) - check_config_value(config, "attention_bias", True) - check_config_value(config, "num_vector_embeds", None) - check_config_value(config, "activation_fn", "gelu-approximate") - check_config_value(config, "num_embeds_ada_norm", 1000) - check_config_value(config, "use_linear_projection", False) - check_config_value(config, "only_cross_attention", False) - check_config_value(config, "cross_attention_norm", True) - check_config_value(config, "double_self_attention", False) - check_config_value(config, "upcast_attention", False) - check_config_value(config, "standardization_norm", "rms_norm") - check_config_value(config, "norm_elementwise_affine", False) - check_config_value(config, "qk_norm", "rms_norm") - check_config_value(config, "positional_embedding_type", "rope") - check_config_value(config, "use_audio_video_cross_attention", True) - check_config_value(config, "share_ff", False) - check_config_value(config, "av_cross_ada_norm", True) - check_config_value(config, "use_middle_indices_grid", True) - check_config_value(config, "num_attention_heads", config.get("audio_num_attention_heads", float("nan"))) - - return LTXModel( - model_type=LTXModelType.AudioVideo, - num_attention_heads=config.get("num_attention_heads", 32), - attention_head_dim=config.get("attention_head_dim", 128), - in_channels=config.get("in_channels", 128), - out_channels=config.get("out_channels", 128), - num_layers=config.get("num_layers", 48), - cross_attention_dim=config.get("cross_attention_dim", 4096), - norm_eps=config.get("norm_eps", 1e-06), - ops=ops, - positional_embedding_theta=config.get("positional_embedding_theta", 10000.0), - positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]), - timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000), - use_middle_indices_grid=config.get("use_middle_indices_grid", True), - audio_num_attention_heads=config.get("audio_num_attention_heads", 32), - audio_attention_head_dim=config.get("audio_attention_head_dim", 64), - audio_in_channels=config.get("audio_in_channels", 128), - audio_out_channels=config.get("audio_out_channels", 128), - audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048), - audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]), - av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1), - rope_type=LTXRopeType(config.get("rope_type", "split")), - double_precision_rope=config.get("frequencies_precision", False) == "float64", - apply_gated_attention=config.get("apply_gated_attention", False), - caption_projection=caption_projection, - audio_caption_projection=audio_caption_projection, - cross_attention_adaln=config.get("cross_attention_adaln", False), - ) - - -class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]): - """ - Configurator for LTX video only model. - Used to create an LTX video only model from a configuration dictionary. - """ - - @classmethod - def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel: - # Build caption projection for 19B model (projection handled in transformer). - caption_projection, _ = _build_caption_projections(config, is_av=False) - - config = config.get("transformer", {}) - - check_config_value(config, "dropout", 0.0) - check_config_value(config, "attention_bias", True) - check_config_value(config, "num_vector_embeds", None) - check_config_value(config, "activation_fn", "gelu-approximate") - check_config_value(config, "num_embeds_ada_norm", 1000) - check_config_value(config, "use_linear_projection", False) - check_config_value(config, "only_cross_attention", False) - check_config_value(config, "cross_attention_norm", True) - check_config_value(config, "double_self_attention", False) - check_config_value(config, "upcast_attention", False) - check_config_value(config, "standardization_norm", "rms_norm") - check_config_value(config, "norm_elementwise_affine", False) - check_config_value(config, "qk_norm", "rms_norm") - check_config_value(config, "positional_embedding_type", "rope") - check_config_value(config, "use_middle_indices_grid", True) - - return LTXModel( - model_type=LTXModelType.VideoOnly, - num_attention_heads=config.get("num_attention_heads", 32), - attention_head_dim=config.get("attention_head_dim", 128), - in_channels=config.get("in_channels", 128), - out_channels=config.get("out_channels", 128), - num_layers=config.get("num_layers", 48), - cross_attention_dim=config.get("cross_attention_dim", 4096), - norm_eps=config.get("norm_eps", 1e-06), - ops=ops, - positional_embedding_theta=config.get("positional_embedding_theta", 10000.0), - positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]), - timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000), - use_middle_indices_grid=config.get("use_middle_indices_grid", True), - rope_type=LTXRopeType(config.get("rope_type", "split")), - double_precision_rope=config.get("frequencies_precision", False) == "float64", - apply_gated_attention=config.get("apply_gated_attention", False), - caption_projection=caption_projection, - cross_attention_adaln=config.get("cross_attention_adaln", False), - ) - - -class LTXAudioOnlyModelConfigurator(ModelConfigurator[LTXModel]): - """ - Configurator for LTX audio only model. - Builds an audio-only LTX model (``model_type=AudioOnly``) so the video - transformer weights are never instantiated or loaded. Useful for - text-to-audio inference where the video branch is unused. - """ - - @classmethod - def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel: - # Build audio caption projection for 19B models (projection handled in transformer). - _, audio_caption_projection = _build_caption_projections(config, is_av=True) - - config = config.get("transformer", {}) - - check_config_value(config, "dropout", 0.0) - check_config_value(config, "attention_bias", True) - check_config_value(config, "num_vector_embeds", None) - check_config_value(config, "activation_fn", "gelu-approximate") - check_config_value(config, "num_embeds_ada_norm", 1000) - check_config_value(config, "use_linear_projection", False) - check_config_value(config, "only_cross_attention", False) - check_config_value(config, "cross_attention_norm", True) - check_config_value(config, "double_self_attention", False) - check_config_value(config, "upcast_attention", False) - check_config_value(config, "standardization_norm", "rms_norm") - check_config_value(config, "norm_elementwise_affine", False) - check_config_value(config, "qk_norm", "rms_norm") - check_config_value(config, "positional_embedding_type", "rope") - check_config_value(config, "use_middle_indices_grid", True) - - return LTXModel( - model_type=LTXModelType.AudioOnly, - num_layers=config.get("num_layers", 48), - norm_eps=config.get("norm_eps", 1e-06), - ops=ops, - timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000), - use_middle_indices_grid=config.get("use_middle_indices_grid", True), - audio_num_attention_heads=config.get("audio_num_attention_heads", 32), - audio_attention_head_dim=config.get("audio_attention_head_dim", 64), - audio_in_channels=config.get("audio_in_channels", 128), - audio_out_channels=config.get("audio_out_channels", 128), - audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048), - audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]), - rope_type=LTXRopeType(config.get("rope_type", "split")), - double_precision_rope=config.get("frequencies_precision", False) == "float64", - apply_gated_attention=config.get("apply_gated_attention", False), - audio_caption_projection=audio_caption_projection, - cross_attention_adaln=config.get("cross_attention_adaln", False), - ) - - -def _build_caption_projections( - config: dict, - is_av: bool, -) -> tuple[torch.nn.Module | None, torch.nn.Module | None]: - """Build caption projections for the transformer when projection is NOT in the text encoder. - 19B models: projection is in the transformer (caption_proj_before_connector=False). - 22B models: projection is in the text encoder, so no projections are created here. - Args: - config: Full model config dict (must contain "transformer" key). - is_av: Whether this is an audio-video model. When False, audio projection is skipped. - Returns: - Tuple of (video_caption_projection, audio_caption_projection), both None for 22B models. - """ - transformer_config = config.get("transformer", {}) - if transformer_config.get("caption_proj_before_connector", False): - return None, None - - with torch.device("meta"): - caption_projection = create_caption_projection(transformer_config) - audio_caption_projection = create_caption_projection(transformer_config, audio=True) if is_av else None - return caption_projection, audio_caption_projection - - -LTXV_MODEL_COMFY_RENAMING_MAP = SDOps("LTXV_MODEL_COMFY_PREFIX_MAP").with_matching(prefix="model.diffusion_model.").with_replacement("model.diffusion_model.", "") - -LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP = ( - SDOps("LTXV_AUDIO_ONLY_MODEL_COMFY_MAP") - .with_matching(prefix="model.diffusion_model.", contains="audio_attn1") - .with_matching(prefix="model.diffusion_model.", contains="audio_attn2") - .with_matching(prefix="model.diffusion_model.", contains="audio_ff") - .with_matching(prefix="model.diffusion_model.", contains="audio_patchify") - .with_matching(prefix="model.diffusion_model.", contains="audio_proj_out") - .with_matching(prefix="model.diffusion_model.", contains="audio_adaln_single") - .with_matching(prefix="model.diffusion_model.", contains="audio_prompt") - .with_matching(prefix="model.diffusion_model.", contains="audio_scale_shift_table") - .with_replacement("model.diffusion_model.", "") -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/ops.py deleted file mode 100644 index ac1acc09f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/ops.py +++ /dev/null @@ -1,105 +0,0 @@ -from typing import List, Protocol - -import torch -from ltx_core.model.transformer.rope import apply_rotary_emb -from ltx_core.utils import rms_norm -from torch import nn - - -class PreAttentionCallable(Protocol): - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - attn_module: nn.Module, - mask: torch.Tensor | None, - pe: torch.Tensor | None, - k_pe: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: ... - - -class PytorchPreAttention(PreAttentionCallable): - def __call__( - self, - q: torch.Tensor, - k: torch.Tensor, - attn_module: nn.Module, - mask: torch.Tensor | None, # noqa: ARG002 - pe: torch.Tensor | None, - k_pe: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - q = attn_module.q_norm(q) - k = attn_module.k_norm(k) - if pe is not None: - q = apply_rotary_emb(q, pe, attn_module.rope_type) - k = apply_rotary_emb(k, pe if k_pe is None else k_pe, attn_module.rope_type) - return q, k - - -class AdaZeroCallable(Protocol): - def __call__( - self, - x: torch.Tensor, - eps: float, - scale: torch.Tensor, - shift: torch.Tensor, - ) -> torch.Tensor: ... - - -class PytorchAdaZeroFunction(AdaZeroCallable): - def __call__( - self, - x: torch.Tensor, - eps: float, - scale: torch.Tensor, - shift: torch.Tensor, - ) -> torch.Tensor: - return rms_norm(x, eps=eps) * (1 + scale) + shift - - -class PostSACallable(Protocol): - def __call__( - self, - x: torch.Tensor, - y: torch.Tensor, - norm_weights: torch.Tensor | None, - eps: float, - gate: torch.Tensor, - ) -> List[torch.Tensor]: ... - - -class PytorchPostSAFunction(PostSACallable): - def __call__( - self, - x: torch.Tensor, - y: torch.Tensor, - norm_weights: torch.Tensor | None, - eps: float, - gate: torch.Tensor, - ) -> List[torch.Tensor]: - x_fma = x + y * gate - return x_fma, rms_norm(x_fma, norm_weights, eps=eps) - - -class GatedAttentionCallable(Protocol): - def __call__( - self, - x: torch.Tensor, - attn_out: torch.Tensor, - attn_module: nn.Module, - ) -> torch.Tensor: ... - - -class PytorchGatedAttention(GatedAttentionCallable): - def __call__( - self, - x: torch.Tensor, - attn_out: torch.Tensor, - attn_module: nn.Module, - ) -> torch.Tensor: - gate_logits = attn_module.to_gate_logits(x) # (B, T, H) - b, t, _ = attn_out.shape - out = attn_out.view(b, t, attn_module.heads, attn_module.dim_head) - gates = 2.0 * torch.sigmoid(gate_logits) # (B, T, H) - out = out * gates.unsqueeze(-1) # (B, T, H, D) * (B, T, H, 1) - return out.view(b, t, attn_module.heads * attn_module.dim_head) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/rope.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/rope.py deleted file mode 100644 index ffc610c48..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/rope.py +++ /dev/null @@ -1,204 +0,0 @@ -import functools -import math -from enum import Enum -from typing import Callable, Tuple - -import numpy as np -import torch -from einops import rearrange - - -class LTXRopeType(Enum): - INTERLEAVED = "interleaved" - SPLIT = "split" - - -def apply_rotary_emb( - input_tensor: torch.Tensor, - freqs_cis: Tuple[torch.Tensor, torch.Tensor], - rope_type: LTXRopeType = LTXRopeType.SPLIT, -) -> torch.Tensor: - if rope_type == LTXRopeType.INTERLEAVED: - # Note: INTERLEAVED rope is a legacy mode. Prefer SPLIT instead. - return apply_interleaved_rotary_emb(input_tensor, *freqs_cis) - elif rope_type == LTXRopeType.SPLIT: - return apply_split_rotary_emb(input_tensor, *freqs_cis) - else: - raise ValueError(f"Invalid rope type: {rope_type}") - - -def apply_interleaved_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor) -> torch.Tensor: - t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) - t1, t2 = t_dup.unbind(dim=-1) - t_dup = torch.stack((-t2, t1), dim=-1) - input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") - - out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs - - return out - - -def apply_split_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor) -> torch.Tensor: - if sin_freqs.shape != cos_freqs.shape: - raise ValueError(f"apply_split_rotary_emb: sin_freqs.shape {tuple(sin_freqs.shape)} must equal cos_freqs.shape {tuple(cos_freqs.shape)}.") - needs_reshape = input_tensor.ndim != 4 and cos_freqs.ndim == 4 - if needs_reshape: - b_freq = cos_freqs.shape[0] - h = cos_freqs.shape[1] - b_in = input_tensor.shape[0] - if b_freq not in (1, b_in): - raise ValueError(f"apply_split_rotary_emb: cos_freqs batch ({b_freq}) must be 1 (broadcast) or equal input_tensor batch ({b_in}).") - # `unflatten` only touches the last dim, keeping the batch and seq dims as - # the input tensor's own symbolic ints under torch.compile. `reshape(b_in, - # t, h, -1)` would have forced Dynamo to specialise those dims because it - # cannot prove `b_in == cos_freqs.shape[0]` and `seq == t` across tensors. - input_tensor = input_tensor.unflatten(-1, (h, -1)).transpose(1, 2) - - split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) - first_half_input = split_input[..., :1, :] - second_half_input = split_input[..., 1:, :] - - output = split_input * cos_freqs.unsqueeze(-2) - first_half_output = output[..., :1, :] - second_half_output = output[..., 1:, :] - - first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) - second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) - - output = rearrange(output, "... d r -> ... (d r)") - if needs_reshape: - # `transpose(1, 2).flatten(-2)` keeps the batch and seq dims symbolic; using - # `reshape(b_in, t, -1)` would force Dynamo to specialise both axes. - output = output.transpose(1, 2).flatten(-2) - - return output - - -@functools.lru_cache(maxsize=5) -def generate_freq_grid_np(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.Tensor: - theta = positional_embedding_theta - start = 1 - end = theta - - n_elem = 2 * positional_embedding_max_pos_count - pow_indices = np.power( - theta, - np.linspace( - np.log(start) / np.log(theta), - np.log(end) / np.log(theta), - inner_dim // n_elem, - dtype=np.float64, - ), - ) - return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32) - - -@functools.lru_cache(maxsize=5) -def generate_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.Tensor: - theta = positional_embedding_theta - start = 1 - end = theta - n_elem = 2 * positional_embedding_max_pos_count - - indices = theta ** ( - torch.linspace( - math.log(start, theta), - math.log(end, theta), - inner_dim // n_elem, - dtype=torch.float32, - ) - ) - indices = indices.to(dtype=torch.float32) - - indices = indices * math.pi / 2 - - return indices - - -def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor: - n_pos_dims = indices_grid.shape[1] - assert n_pos_dims == len(max_pos), f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})" - fractional_positions = torch.stack( - [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)], - dim=-1, - ) - return fractional_positions - - -def generate_freqs(indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool) -> torch.Tensor: - if use_middle_indices_grid: - assert len(indices_grid.shape) == 4 - assert indices_grid.shape[-1] == 2 - indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1] - indices_grid = (indices_grid_start + indices_grid_end) / 2.0 - elif len(indices_grid.shape) == 4: - indices_grid = indices_grid[..., 0] - - fractional_positions = get_fractional_positions(indices_grid, max_pos) - indices = indices.to(device=fractional_positions.device) - - freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2) - return freqs - - -def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]: - cos_freq = freqs.cos() - sin_freq = freqs.sin() - - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) - - cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) - sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) - - # Reshape freqs to be compatible with multi-head attention - b = cos_freq.shape[0] - t = cos_freq.shape[1] - - cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1) - sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1) - - cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) - sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) - return cos_freq, sin_freq - - -def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]: - cos_freq = freqs.cos().repeat_interleave(2, dim=-1) - sin_freq = freqs.sin().repeat_interleave(2, dim=-1) - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size]) - cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) - sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) - return cos_freq, sin_freq - - -def precompute_freqs_cis( - indices_grid: torch.Tensor, - dim: int, - out_dtype: torch.dtype, - theta: float = 10000.0, - max_pos: list[int] | None = None, - use_middle_indices_grid: bool = False, - num_attention_heads: int = 32, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch, -) -> tuple[torch.Tensor, torch.Tensor]: - if max_pos is None: - max_pos = [20, 2048, 2048] - - indices = freq_grid_generator(theta, indices_grid.shape[1], dim) - freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid) - - if rope_type == LTXRopeType.SPLIT: - expected_freqs = dim // 2 - current_freqs = freqs.shape[-1] - pad_size = expected_freqs - current_freqs - cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads) - else: - # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only - n_elem = 2 * indices_grid.shape[1] - cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem) - return cos_freq.to(out_dtype), sin_freq.to(out_dtype) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/text_projection.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/text_projection.py deleted file mode 100644 index 5046755d2..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/text_projection.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch - - -class PixArtAlphaTextProjection(torch.nn.Module): - """ - Projects caption embeddings using dual linear layers. - Flow: linear_1 → activation → linear_2 - Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py - """ - - def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): - super().__init__() - if out_features is None: - out_features = hidden_size - self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) - if act_fn == "gelu_tanh": - self.act_1 = torch.nn.GELU(approximate="tanh") - elif act_fn == "silu": - self.act_1 = torch.nn.SiLU() - else: - raise ValueError(f"Unknown activation function: {act_fn}") - self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) - - def forward(self, caption: torch.Tensor) -> torch.Tensor: - hidden_states = self.linear_1(caption) - hidden_states = self.act_1(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states - - -def create_caption_projection(transformer_config: dict, audio: bool = False) -> PixArtAlphaTextProjection: - """Create a caption projection for the transformer (V1/19B only).""" - caption_channels = transformer_config["caption_channels"] - if audio: - inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"] - else: - inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"] - return PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/timestep_embedding.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/timestep_embedding.py deleted file mode 100644 index 820a95589..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/timestep_embedding.py +++ /dev/null @@ -1,143 +0,0 @@ -import math - -import torch - - -def get_timestep_embedding( - timesteps: torch.Tensor, - embedding_dim: int, - flip_sin_to_cos: bool = False, - downscale_freq_shift: float = 1, - scale: float = 1, - max_period: int = 10000, -) -> torch.Tensor: - """ - This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. - Args - timesteps (torch.Tensor): - a 1-D Tensor of N indices, one per batch element. These may be fractional. - embedding_dim (int): - the dimension of the output. - flip_sin_to_cos (bool): - Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) - downscale_freq_shift (float): - Controls the delta between frequencies between dimensions - scale (float): - Scaling factor applied to the embeddings. - max_period (int): - Controls the maximum frequency of the embeddings - Returns - torch.Tensor: an [N x dim] Tensor of positional embeddings. - """ - assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" - - half_dim = embedding_dim // 2 - exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) - exponent = exponent / (half_dim - downscale_freq_shift) - - emb = torch.exp(exponent) - emb = timesteps[:, None].float() * emb[None, :] - - # scale embeddings - emb = scale * emb - - # concat sine and cosine embeddings - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) - - # flip sine and cosine embeddings - if flip_sin_to_cos: - emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) - - # zero pad - if embedding_dim % 2 == 1: - emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) - return emb - - -class TimestepEmbedding(torch.nn.Module): - def __init__( - self, - in_channels: int, - time_embed_dim: int, - out_dim: int | None = None, - post_act_fn: str | None = None, - cond_proj_dim: int | None = None, - sample_proj_bias: bool = True, - ): - super().__init__() - - self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) - - if cond_proj_dim is not None: - self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) - else: - self.cond_proj = None - - self.act = torch.nn.SiLU() - time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim - - self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) - - if post_act_fn is None: - self.post_act = None - - def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: - if condition is not None: - sample = sample + self.cond_proj(condition) - sample = self.linear_1(sample) - - if self.act is not None: - sample = self.act(sample) - - sample = self.linear_2(sample) - - if self.post_act is not None: - sample = self.post_act(sample) - return sample - - -class Timesteps(torch.nn.Module): - def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): - super().__init__() - self.num_channels = num_channels - self.flip_sin_to_cos = flip_sin_to_cos - self.downscale_freq_shift = downscale_freq_shift - self.scale = scale - - def forward(self, timesteps: torch.Tensor) -> torch.Tensor: - t_emb = get_timestep_embedding( - timesteps, - self.num_channels, - flip_sin_to_cos=self.flip_sin_to_cos, - downscale_freq_shift=self.downscale_freq_shift, - scale=self.scale, - ) - return t_emb - - -class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): - """ - For PixArt-Alpha. - Reference: - https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 - """ - - def __init__( - self, - embedding_dim: int, - size_emb_dim: int, - ): - super().__init__() - - self.outdim = size_emb_dim - self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) - - def forward( - self, - timestep: torch.Tensor, - hidden_dtype: torch.dtype, - ) -> torch.Tensor: - timesteps_proj = self.time_proj(timestep) - timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) - return timesteps_emb diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer.py deleted file mode 100644 index e50e7e1fc..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer.py +++ /dev/null @@ -1,444 +0,0 @@ -from dataclasses import dataclass, field, replace - -import torch -from ltx_core.model.transformer.adaln import adaln_embedding_coefficient -from ltx_core.model.transformer.attention import ( - Attention, - AttentionCallable, - AttentionFunction, - AttentionOps, - MaskedAttentionCallable, - MaskedAttentionFunction, -) -from ltx_core.model.transformer.feed_forward import FeedForward -from ltx_core.model.transformer.ops import ( - AdaZeroCallable, - GatedAttentionCallable, - PostSACallable, - PreAttentionCallable, - PytorchAdaZeroFunction, - PytorchGatedAttention, - PytorchPostSAFunction, - PytorchPreAttention, -) -from ltx_core.model.transformer.rope import LTXRopeType -from ltx_core.model.transformer.transformer_args import TransformerArgs - - -@dataclass -class TransformerConfig: - dim: int - heads: int - d_head: int - context_dim: int - apply_gated_attention: bool = False - cross_attention_adaln: bool = False - - -@dataclass(frozen=True) -class TransformerOpsConfig: - """Pluggable ops for :class:`BasicAVTransformerBlock`. - Use :meth:`from_functions` to construct from enum values or partial overrides - without spelling out a full :class:`AttentionOps`. - """ - - attention_ops: AttentionOps = field(default_factory=AttentionOps) - ada_zero_function: AdaZeroCallable = field(default_factory=PytorchAdaZeroFunction) - post_sa_function: PostSACallable = field(default_factory=PytorchPostSAFunction) - - @classmethod - def from_functions( - cls, - attention: AttentionFunction | AttentionCallable = AttentionFunction.AUTOMATIC, - masked_attention: MaskedAttentionFunction | MaskedAttentionCallable = MaskedAttentionFunction.AUTOMATIC, - preattention: PreAttentionCallable | None = None, - gated_attention: GatedAttentionCallable | None = None, - ada_zero: AdaZeroCallable | None = None, - post_sa: PostSACallable | None = None, - ) -> "TransformerOpsConfig": - """Build a config from individual functions or enums. Each *None* slot - falls back to the standard PyTorch implementation.""" - attention_callable = attention.to_callable() if isinstance(attention, AttentionFunction) else attention - masked_callable = masked_attention.to_callable() if isinstance(masked_attention, MaskedAttentionFunction) else masked_attention - attention_ops = AttentionOps( - attention_function=attention_callable, - masked_attention_function=masked_callable, - preattention_function=preattention if preattention is not None else PytorchPreAttention(), - gated_attention_function=(gated_attention if gated_attention is not None else PytorchGatedAttention()), - ) - return cls( - attention_ops=attention_ops, - ada_zero_function=ada_zero if ada_zero is not None else PytorchAdaZeroFunction(), - post_sa_function=post_sa if post_sa is not None else PytorchPostSAFunction(), - ) - - -# Frozen, so safe to share as a default argument across callers that want the -# stock PyTorch ops without explicit construction. -DEFAULT_TRANSFORMER_OPS = TransformerOpsConfig() - - -class BasicAVTransformerBlock(torch.nn.Module): - def __init__( - self, - video: TransformerConfig | None = None, - audio: TransformerConfig | None = None, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - norm_eps: float = 1e-6, - ops: TransformerOpsConfig | None = None, - ): - super().__init__() - - if ops is None: - ops = TransformerOpsConfig() - self.ada_zero_function = ops.ada_zero_function - self.post_sa_function = ops.post_sa_function - if video is not None: - self.attn1 = Attention( - query_dim=video.dim, - heads=video.heads, - dim_head=video.d_head, - context_dim=None, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=video.apply_gated_attention, - ) - self.attn2 = Attention( - query_dim=video.dim, - context_dim=video.context_dim, - heads=video.heads, - dim_head=video.d_head, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=video.apply_gated_attention, - ) - self.ff = FeedForward(video.dim, dim_out=video.dim) - video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln) - self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim)) - - if audio is not None: - self.audio_attn1 = Attention( - query_dim=audio.dim, - heads=audio.heads, - dim_head=audio.d_head, - context_dim=None, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=audio.apply_gated_attention, - ) - self.audio_attn2 = Attention( - query_dim=audio.dim, - context_dim=audio.context_dim, - heads=audio.heads, - dim_head=audio.d_head, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=audio.apply_gated_attention, - ) - self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim) - audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln) - self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim)) - - if audio is not None and video is not None: - # Q: Video, K,V: Audio - self.audio_to_video_attn = Attention( - query_dim=video.dim, - context_dim=audio.dim, - heads=audio.heads, - dim_head=audio.d_head, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=video.apply_gated_attention, - ) - - # Q: Audio, K,V: Video - self.video_to_audio_attn = Attention( - query_dim=audio.dim, - context_dim=video.dim, - heads=audio.heads, - dim_head=audio.d_head, - rope_type=rope_type, - norm_eps=norm_eps, - ops=ops.attention_ops, - apply_gated_attention=audio.apply_gated_attention, - ) - - self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim)) - self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim)) - - self.cross_attention_adaln = (video is not None and video.cross_attention_adaln) or (audio is not None and audio.cross_attention_adaln) - - if self.cross_attention_adaln and video is not None: - self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim)) - if self.cross_attention_adaln and audio is not None: - self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.dim)) - - self.norm_eps = norm_eps - - def get_ada_values(self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice) -> tuple[torch.Tensor, ...]: - num_ada_params = scale_shift_table.shape[0] - - ada_values = ( - scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) - + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] - ).unbind(dim=2) - return ada_values - - def get_av_ca_ada_values( - self, - scale_shift_table: torch.Tensor, - batch_size: int, - scale_shift_timestep: torch.Tensor, - gate_timestep: torch.Tensor, - scale_shift_indices: slice, - num_scale_shift_values: int = 4, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - scale_shift_ada_values = self.get_ada_values(scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, scale_shift_indices) - gate_ada_values = self.get_ada_values(scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)) - - scale, shift = (t.squeeze(2) for t in scale_shift_ada_values) - (gate,) = (t.squeeze(2) for t in gate_ada_values) - - return scale, shift, gate - - def _apply_text_cross_attention( - self, - x_normed: torch.Tensor, - context: torch.Tensor, - attn: AttentionCallable, - scale_shift_table: torch.Tensor, - prompt_scale_shift_table: torch.Tensor | None, - timestep: torch.Tensor, - prompt_timestep: torch.Tensor | None, - context_mask: torch.Tensor | None, - cross_attention_adaln: bool = False, - ) -> torch.Tensor: - """Apply text cross-attention, with optional AdaLN modulation. - ``x_normed`` is the RMS-normalized self-attention output produced by - ``post_sa_function`` -- this method does not normalize again. - """ - if cross_attention_adaln: - shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x_normed.shape[0], timestep, slice(6, 9)) - return apply_cross_attention_adaln( - x_normed, - context, - attn, - shift_q, - scale_q, - gate, - prompt_scale_shift_table, - prompt_timestep, - context_mask, - ) - return attn(x_normed, context=context, mask=context_mask) - - def forward( # noqa: PLR0915 - self, - video: TransformerArgs | None, - audio: TransformerArgs | None, - video_kv_cache: dict | None = None, - audio_kv_cache: dict | None = None, - video_current_start: int = 0, - audio_current_start: int = 0, - defer_cache_update: bool = False, - detach_cache_update: bool = False, - ) -> tuple[TransformerArgs | None, TransformerArgs | None]: - if video is None and audio is None: - raise ValueError("At least one of video or audio must be provided") - - vx = video.x if video is not None else None - ax = audio.x if audio is not None else None - video_cache_update_info = None - audio_cache_update_info = None - - run_vx = video is not None and video.enabled and vx.numel() > 0 - run_ax = audio is not None and audio.enabled and ax.numel() > 0 - - run_a2v = run_vx and (audio is not None and ax.numel() > 0) - run_v2a = run_ax and (video is not None and vx.numel() > 0) - - if run_vx: - vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)) - norm_vx = self.ada_zero_function(vx, self.norm_eps, vscale_msa, vshift_msa) - del vshift_msa, vscale_msa - - vx_msa_out = self.attn1( - norm_vx, - pe=video.positional_embeddings, - mask=video.self_attention_mask, - perturbation_mask=video.self_attn_perturbation_mask, - all_perturbed=video.self_attn_all_perturbed, - kv_cache=video_kv_cache, - current_start=video_current_start, - defer_cache_update=defer_cache_update, - detach_cache_update=detach_cache_update, - ) - if isinstance(vx_msa_out, tuple): - vx_msa_out, video_cache_update_info = vx_msa_out - vx, vx_normed = self.post_sa_function(vx, vx_msa_out, None, self.norm_eps, vgate_msa) - del vgate_msa, norm_vx, vx_msa_out - vx = vx + self._apply_text_cross_attention( - vx_normed, - video.context, - self.attn2, - self.scale_shift_table, - getattr(self, "prompt_scale_shift_table", None), - video.timesteps, - video.prompt_timestep, - video.context_mask, - cross_attention_adaln=self.cross_attention_adaln, - ) - del vx_normed - - if run_ax: - ashift_msa, ascale_msa, agate_msa = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)) - - norm_ax = self.ada_zero_function(ax, self.norm_eps, ascale_msa, ashift_msa) - del ashift_msa, ascale_msa - ax_msa_out = self.audio_attn1( - norm_ax, - pe=audio.positional_embeddings, - mask=audio.self_attention_mask, - perturbation_mask=audio.self_attn_perturbation_mask, - all_perturbed=audio.self_attn_all_perturbed, - kv_cache=audio_kv_cache, - current_start=audio_current_start, - defer_cache_update=defer_cache_update, - detach_cache_update=detach_cache_update, - ) - if isinstance(ax_msa_out, tuple): - ax_msa_out, audio_cache_update_info = ax_msa_out - ax, ax_normed = self.post_sa_function(ax, ax_msa_out, None, self.norm_eps, agate_msa) - del agate_msa, norm_ax, ax_msa_out - ax = ax + self._apply_text_cross_attention( - ax_normed, - audio.context, - self.audio_attn2, - self.audio_scale_shift_table, - getattr(self, "audio_prompt_scale_shift_table", None), - audio.timesteps, - audio.prompt_timestep, - audio.context_mask, - cross_attention_adaln=self.cross_attention_adaln, - ) - del ax_normed - - # Audio - Video cross attention. - if run_a2v or run_v2a: - # Snapshot vx/ax before A2V mutates vx; V2A's video keys/values must - # use the pre-A2V state so direction order doesn't bias the result. - vx_pre_av = vx - ax_pre_av = ax - if run_a2v and not video.cross_attn_skip_all: - scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_video, - vx.shape[0], - video.cross_scale_shift_timestep, - video.cross_gate_timestep, - slice(0, 2), - ) - a2v_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_a2v, shift_ca_video_a2v) - del scale_ca_video_a2v, shift_ca_video_a2v - - scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_audio, - ax.shape[0], - audio.cross_scale_shift_timestep, - audio.cross_gate_timestep, - slice(0, 2), - ) - a2v_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_a2v, shift_ca_audio_a2v) - del scale_ca_audio_a2v, shift_ca_audio_a2v - vx = vx + ( - self.audio_to_video_attn( - a2v_vx_scaled, - context=a2v_ax_scaled, - pe=video.cross_positional_embeddings, - k_pe=audio.cross_positional_embeddings, - mask=video.cross_attention_mask, - ) - * gate_out_a2v - * video.cross_attn_perturbation_mask - ) - del gate_out_a2v, a2v_vx_scaled, a2v_ax_scaled - - if run_v2a and not audio.cross_attn_skip_all: - scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_audio, - ax.shape[0], - audio.cross_scale_shift_timestep, - audio.cross_gate_timestep, - slice(2, 4), - ) - v2a_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_v2a, shift_ca_audio_v2a) - del scale_ca_audio_v2a, shift_ca_audio_v2a - scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values( - self.scale_shift_table_a2v_ca_video, - vx.shape[0], - video.cross_scale_shift_timestep, - video.cross_gate_timestep, - slice(2, 4), - ) - v2a_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_v2a, shift_ca_video_v2a) - del scale_ca_video_v2a, shift_ca_video_v2a - ax = ax + ( - self.video_to_audio_attn( - v2a_ax_scaled, - context=v2a_vx_scaled, - pe=audio.cross_positional_embeddings, - k_pe=video.cross_positional_embeddings, - mask=audio.cross_attention_mask, - ) - * gate_out_v2a - * audio.cross_attn_perturbation_mask - ) - del gate_out_v2a, v2a_vx_scaled, v2a_ax_scaled - del vx_pre_av, ax_pre_av - - if run_vx: - vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)) - vx_scaled = self.ada_zero_function(vx, self.norm_eps, vscale_mlp, vshift_mlp) - vx = vx + self.ff(vx_scaled) * vgate_mlp - - del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled - - if run_ax: - ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)) - ax_scaled = self.ada_zero_function(ax, self.norm_eps, ascale_mlp, ashift_mlp) - ax = ax + self.audio_ff(ax_scaled) * agate_mlp - - del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled - - video_out = replace(video, x=vx) if video is not None else None - audio_out = replace(audio, x=ax) if audio is not None else None - if video_cache_update_info is not None or audio_cache_update_info is not None: - return video_out, audio_out, {"video": video_cache_update_info, "audio": audio_cache_update_info} - return video_out, audio_out - - -def apply_cross_attention_adaln( - x_normed: torch.Tensor, - context: torch.Tensor, - attn: AttentionCallable, - q_shift: torch.Tensor, - q_scale: torch.Tensor, - q_gate: torch.Tensor, - prompt_scale_shift_table: torch.Tensor, - prompt_timestep: torch.Tensor, - context_mask: torch.Tensor | None = None, -) -> torch.Tensor: - """Apply query/key AdaLN modulation then cross-attention. - ``x_normed`` is already RMS-normalized by ``post_sa_function``; this only - applies the affine (scale/shift) modulation, so the normalization is not - repeated here. - """ - batch_size = x_normed.shape[0] - shift_kv, scale_kv = (prompt_scale_shift_table[None, None].to(device=x_normed.device, dtype=x_normed.dtype) + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)).unbind(dim=2) - attn_input = x_normed * (1 + q_scale) + q_shift - encoder_hidden_states = context * (1 + scale_kv) + shift_kv - return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer_args.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer_args.py deleted file mode 100644 index c9b626091..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer_args.py +++ /dev/null @@ -1,351 +0,0 @@ -from dataclasses import dataclass, replace - -import torch -from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType -from ltx_core.model.transformer.adaln import AdaLayerNormSingle -from ltx_core.model.transformer.modality import Modality -from ltx_core.model.transformer.rope import ( - LTXRopeType, - generate_freq_grid_np, - generate_freq_grid_pytorch, - precompute_freqs_cis, -) -from torch.nn.attention.flex_attention import BlockMask -from torch.utils import _pytree - - -@dataclass(frozen=True) -class TransformerArgs: - x: torch.Tensor - context: torch.Tensor - context_mask: torch.Tensor - timesteps: torch.Tensor - embedded_timestep: torch.Tensor - positional_embeddings: tuple[torch.Tensor, torch.Tensor] - cross_positional_embeddings: tuple[torch.Tensor, torch.Tensor] | None - cross_scale_shift_timestep: torch.Tensor | None - cross_gate_timestep: torch.Tensor | None - enabled: bool - prompt_timestep: torch.Tensor | None = None - self_attention_mask: torch.Tensor | BlockMask | None = None - cross_attention_mask: torch.Tensor | BlockMask | None = None - # Per-block perturbation state, precomputed by `LTXModel._process_transformer_blocks` - # so the block forward needs no per-block identity. The bool shortcuts - # (`*_all_perturbed`, `cross_attn_skip_all`) are Python bools that Dynamo specialises - # on — fine because they're stable across denoising steps for a fixed perturbation - # config. - self_attn_perturbation_mask: torch.Tensor | None = None - self_attn_all_perturbed: bool = False - cross_attn_perturbation_mask: torch.Tensor | None = None - cross_attn_skip_all: bool = False - - -# FSDP2 locates tensors in module outputs through pytree flattening before -# registering its pre-backward all-gather hooks. -_pytree.register_dataclass(TransformerArgs) - - -class BlockPerturbationsProcessor: - """Per-block preparation of ``TransformerArgs``. - The base implementation returns a copy of ``args`` with this block's - precomputed perturbation flags and masks attached. Subclasses can layer in - operations that must run on each block's inputs but stay outside the - compile boundary -- e.g. ``torch._dynamo.mark_dynamic`` for - shape-polymorphic block compilation (see ``compiling.py``). Swapping the - processor on an ``LTXModel`` instance is how compile transforms opt in to - such behaviour without baking it into the model's forward. - ``self_attn_perturbation_mask`` is None when all or none of the batch is - perturbed (the attention call can take the shortcut path). ``cross_attn_*`` - is None when every sample skips the cross-attention entirely. - """ - - def __call__( - self, - args: "TransformerArgs", - perturbations: BatchedPerturbationConfig, - block_idx: int, - self_attn_type: PerturbationType, - cross_attn_type: PerturbationType, - ) -> "TransformerArgs": - device, dtype = args.x.device, args.x.dtype - - all_self = perturbations.all_in_batch(self_attn_type, block_idx) - any_self = perturbations.any_in_batch(self_attn_type, block_idx) - self_mask: torch.Tensor | None = None - if any_self and not all_self: - self_mask = perturbations.mask(self_attn_type, block_idx, device, dtype).view(-1, 1, 1) - - all_cross = perturbations.all_in_batch(cross_attn_type, block_idx) - cross_mask: torch.Tensor | None = None - if not all_cross: - cross_mask = perturbations.mask(cross_attn_type, block_idx, device, dtype).view(-1, 1, 1) - - return replace( - args, - self_attn_perturbation_mask=self_mask, - self_attn_all_perturbed=all_self, - cross_attn_perturbation_mask=cross_mask, - cross_attn_skip_all=all_cross, - ) - - -class TransformerArgsPreprocessor: - def __init__( # noqa: PLR0913 - self, - patchify_proj: torch.nn.Linear, - adaln: AdaLayerNormSingle, - inner_dim: int, - max_pos: list[int], - num_attention_heads: int, - use_middle_indices_grid: bool, - timestep_scale_multiplier: int, - double_precision_rope: bool, - positional_embedding_theta: float, - rope_type: LTXRopeType, - caption_projection: torch.nn.Module | None = None, - prompt_adaln: AdaLayerNormSingle | None = None, - ) -> None: - self.patchify_proj = patchify_proj - self.adaln = adaln - self.inner_dim = inner_dim - self.max_pos = max_pos - self.num_attention_heads = num_attention_heads - self.use_middle_indices_grid = use_middle_indices_grid - self.timestep_scale_multiplier = timestep_scale_multiplier - self.double_precision_rope = double_precision_rope - self.positional_embedding_theta = positional_embedding_theta - self.rope_type = rope_type - self.caption_projection = caption_projection - self.prompt_adaln = prompt_adaln - - def _prepare_timestep(self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare timestep embeddings.""" - timestep_scaled = timestep * self.timestep_scale_multiplier - timestep, embedded_timestep = adaln( - timestep_scaled.flatten(), - hidden_dtype=hidden_dtype, - ) - # Second dimension is 1 or number of tokens (if timestep_per_token) - timestep = timestep.view(batch_size, -1, timestep.shape[-1]) - embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) - - return timestep, embedded_timestep - - def _prepare_context( - self, - context: torch.Tensor, - x: torch.Tensor, - ) -> torch.Tensor: - """Prepare context for transformer blocks.""" - if self.caption_projection is not None: - context = self.caption_projection(context) - batch_size = x.shape[0] - return context.view(batch_size, -1, x.shape[-1]) - - def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: - """Prepare attention mask.""" - if attention_mask is None or torch.is_floating_point(attention_mask): - return attention_mask - - return (attention_mask - 1).to(x_dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(x_dtype).max - - def _prepare_self_attention_mask(self, attention_mask: torch.Tensor | BlockMask | None, x_dtype: torch.dtype) -> torch.Tensor | BlockMask | None: - """Prepare self-attention mask by converting [0,1] values to additive log-space bias. - Input shape: 3D ``(B, T_q, T_k)`` with values in [0, 1]. The dense form - is ``(B, T, T)``; broadcastable forms like ``(1, 1, T)`` (key-only - padding) or ``(B, 1, T)`` are also valid and yield a correspondingly - broadcastable output. - Output shape: ``(B, 1, T_q, T_k)`` (heads dim inserted) with 0.0 for - full attention and a large negative value for masked positions. - Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum - representable value). Strictly positive entries are converted via log-space for - smooth attenuation, with small values clamped for numerical stability. - Returns None if input is None (no masking). - """ - if attention_mask is None or isinstance(attention_mask, BlockMask): - return attention_mask - - # Convert [0, 1] attention mask to additive log-space bias: - # 1.0 -> log(1.0) = 0.0 (no bias, full attention) - # 0.0 -> finfo.min (fully masked) - finfo = torch.finfo(x_dtype) - eps = finfo.tiny - - bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype) - positive = attention_mask > 0 - if positive.any(): - bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype) - - return bias.unsqueeze(1) # (B, 1, T_q, T_k) for head broadcast - - def _prepare_positional_embeddings( - self, - positions: torch.Tensor, - inner_dim: int, - max_pos: list[int], - use_middle_indices_grid: bool, - num_attention_heads: int, - x_dtype: torch.dtype, - ) -> torch.Tensor: - """Prepare positional embeddings.""" - freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch - pe = precompute_freqs_cis( - positions, - dim=inner_dim, - out_dtype=x_dtype, - theta=self.positional_embedding_theta, - max_pos=max_pos, - use_middle_indices_grid=use_middle_indices_grid, - num_attention_heads=num_attention_heads, - rope_type=self.rope_type, - freq_grid_generator=freq_grid_generator, - ) - return pe - - def prepare( - self, - modality: Modality, - cross_modality: Modality | None = None, # noqa: ARG002 - ) -> TransformerArgs: - x = self.patchify_proj(modality.latent) - batch_size = x.shape[0] - timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, self.adaln, batch_size, modality.latent.dtype) - prompt_timestep = None - if self.prompt_adaln is not None: - prompt_timestep, _ = self._prepare_timestep(modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype) - context = self._prepare_context(modality.context, x) - attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype) - pe = self._prepare_positional_embeddings( - positions=modality.positions, - inner_dim=self.inner_dim, - max_pos=self.max_pos, - use_middle_indices_grid=self.use_middle_indices_grid, - num_attention_heads=self.num_attention_heads, - x_dtype=modality.latent.dtype, - ) - self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype) - cross_attention_mask = self._prepare_self_attention_mask( - getattr(modality, "cross_attention_mask", None), - modality.latent.dtype, - ) - return TransformerArgs( - x=x, - context=context, - context_mask=attention_mask, - timesteps=timestep, - embedded_timestep=embedded_timestep, - positional_embeddings=pe, - cross_positional_embeddings=None, - cross_scale_shift_timestep=None, - cross_gate_timestep=None, - enabled=modality.enabled, - prompt_timestep=prompt_timestep, - self_attention_mask=self_attention_mask, - cross_attention_mask=cross_attention_mask, - ) - - -class MultiModalTransformerArgsPreprocessor: - def __init__( # noqa: PLR0913 - self, - patchify_proj: torch.nn.Linear, - adaln: AdaLayerNormSingle, - cross_scale_shift_adaln: AdaLayerNormSingle, - cross_gate_adaln: AdaLayerNormSingle, - inner_dim: int, - max_pos: list[int], - num_attention_heads: int, - cross_pe_max_pos: int, - use_middle_indices_grid: bool, - audio_cross_attention_dim: int, - timestep_scale_multiplier: int, - double_precision_rope: bool, - positional_embedding_theta: float, - rope_type: LTXRopeType, - av_ca_timestep_scale_multiplier: int, - caption_projection: torch.nn.Module | None = None, - prompt_adaln: AdaLayerNormSingle | None = None, - ) -> None: - self.simple_preprocessor = TransformerArgsPreprocessor( - patchify_proj=patchify_proj, - adaln=adaln, - inner_dim=inner_dim, - max_pos=max_pos, - num_attention_heads=num_attention_heads, - use_middle_indices_grid=use_middle_indices_grid, - timestep_scale_multiplier=timestep_scale_multiplier, - double_precision_rope=double_precision_rope, - positional_embedding_theta=positional_embedding_theta, - rope_type=rope_type, - caption_projection=caption_projection, - prompt_adaln=prompt_adaln, - ) - self.cross_scale_shift_adaln = cross_scale_shift_adaln - self.cross_gate_adaln = cross_gate_adaln - self.cross_pe_max_pos = cross_pe_max_pos - self.audio_cross_attention_dim = audio_cross_attention_dim - self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier - - def prepare( - self, - modality: Modality, - cross_modality: Modality | None = None, - ) -> TransformerArgs: - transformer_args = self.simple_preprocessor.prepare(modality) - if cross_modality is None: - return transformer_args - - if cross_modality.sigma.numel() > 1: - if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]: - raise ValueError("Cross modality sigma must have the same batch size as the modality") - if cross_modality.sigma.ndim != 1: - raise ValueError("Cross modality sigma must be a 1D tensor") - - cross_pe = self.simple_preprocessor._prepare_positional_embeddings( - positions=modality.positions[:, 0:1, :], - inner_dim=self.audio_cross_attention_dim, - max_pos=[self.cross_pe_max_pos], - use_middle_indices_grid=True, - num_attention_heads=self.simple_preprocessor.num_attention_heads, - x_dtype=modality.latent.dtype, - ) - - cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( - modality_timesteps=modality.timesteps, - cross_modality_sigma=cross_modality.sigma, - timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, - batch_size=transformer_args.x.shape[0], - hidden_dtype=modality.latent.dtype, - ) - - return replace( - transformer_args, - cross_positional_embeddings=cross_pe, - cross_scale_shift_timestep=cross_scale_shift_timestep, - cross_gate_timestep=cross_gate_timestep, - ) - - def _prepare_cross_attention_timestep( - self, - modality_timesteps: torch.Tensor, - cross_modality_sigma: torch.Tensor, - timestep_scale_multiplier: int, - batch_size: int, - hidden_dtype: torch.dtype, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare A-V cross-attention AdaLN inputs.""" - av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier - - scale_shift_timestep, _ = self.cross_scale_shift_adaln( - (modality_timesteps * timestep_scale_multiplier).flatten(), - hidden_dtype=hidden_dtype, - ) - scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) - - gate_noise_timestep, _ = self.cross_gate_adaln( - (cross_modality_sigma * timestep_scale_multiplier * av_ca_factor).flatten(), - hidden_dtype=hidden_dtype, - ) - gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) - - return scale_shift_timestep, gate_noise_timestep diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/__init__.py deleted file mode 100644 index 4836bfc9f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Latent upsampler model components.""" - -from ltx_core.model.upsampler.model import LatentUpsampler, upsample_video -from ltx_core.model.upsampler.model_configurator import LatentUpsamplerConfigurator - -__all__ = [ - "LatentUpsampler", - "LatentUpsamplerConfigurator", - "upsample_video", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/blur_downsample.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/blur_downsample.py deleted file mode 100644 index ccc014973..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/blur_downsample.py +++ /dev/null @@ -1,53 +0,0 @@ -import math - -import torch -import torch.nn.functional as F -from einops import rearrange - - -class BlurDownsample(torch.nn.Module): - """ - Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. - Applies only on H,W. Works for dims=2 or dims=3 (per-frame). - """ - - def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: - super().__init__() - assert dims in (2, 3) - assert isinstance(stride, int) - assert stride >= 1 - assert kernel_size >= 3 - assert kernel_size % 2 == 1 - self.dims = dims - self.stride = stride - self.kernel_size = kernel_size - - # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from - # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and - # provides a smooth approximation of a Gaussian filter (often called a "binomial filter"). - # The 2D kernel is constructed as the outer product and normalized. - k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)]) - k2d = k[:, None] @ k[None, :] - k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size) - self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.stride == 1: - return x - - if self.dims == 2: - return self._apply_2d(x) - else: - # dims == 3: apply per-frame on H,W - b, _, f, _, _ = x.shape - x = rearrange(x, "b c f h w -> (b f) c h w") - x = self._apply_2d(x) - h2, w2 = x.shape[-2:] - x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2) - return x - - def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor: - c = x2d.shape[1] - weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise - x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) - return x2d diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model.py deleted file mode 100644 index 7d16eb6d5..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model.py +++ /dev/null @@ -1,139 +0,0 @@ -import torch -from einops import rearrange -from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND -from ltx_core.model.upsampler.res_block import ResBlock -from ltx_core.model.upsampler.spatial_rational_resampler import SpatialRationalResampler -from ltx_core.model.video_vae import VideoEncoder - - -class LatentUpsampler(torch.nn.Module): - """ - Model to upsample VAE latents spatially and/or temporally. - Args: - in_channels (`int`): Number of channels in the input latent - mid_channels (`int`): Number of channels in the middle layers - num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling) - dims (`int`): Number of dimensions for convolutions (2 or 3) - spatial_upsample (`bool`): Whether to spatially upsample the latent - temporal_upsample (`bool`): Whether to temporally upsample the latent - spatial_scale (`float`): Scale factor for spatial upsampling - rational_resampler (`bool`): Whether to use a rational resampler for spatial upsampling - """ - - def __init__( - self, - in_channels: int = 128, - mid_channels: int = 512, - num_blocks_per_stage: int = 4, - dims: int = 3, - spatial_upsample: bool = True, - temporal_upsample: bool = False, - spatial_scale: float = 2.0, - rational_resampler: bool = False, - ): - super().__init__() - - self.in_channels = in_channels - self.mid_channels = mid_channels - self.num_blocks_per_stage = num_blocks_per_stage - self.dims = dims - self.spatial_upsample = spatial_upsample - self.temporal_upsample = temporal_upsample - self.spatial_scale = float(spatial_scale) - self.rational_resampler = rational_resampler - - conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1) - self.initial_norm = torch.nn.GroupNorm(32, mid_channels) - self.initial_activation = torch.nn.SiLU() - - self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) - - if spatial_upsample and temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(3), - ) - elif spatial_upsample: - if rational_resampler: - self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale) - else: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(2), - ) - elif temporal_upsample: - self.upsampler = torch.nn.Sequential( - torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), - PixelShuffleND(1), - ) - else: - raise ValueError("Either spatial_upsample or temporal_upsample must be True") - - self.post_upsample_res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]) - - self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1) - - def forward(self, latent: torch.Tensor) -> torch.Tensor: - b, _, f, _, _ = latent.shape - - if self.dims == 2: - x = rearrange(latent, "b c f h w -> (b f) c h w") - x = self.initial_conv(x) - x = self.initial_norm(x) - x = self.initial_activation(x) - - for block in self.res_blocks: - x = block(x) - - x = self.upsampler(x) - - for block in self.post_upsample_res_blocks: - x = block(x) - - x = self.final_conv(x) - x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) - else: - x = self.initial_conv(latent) - x = self.initial_norm(x) - x = self.initial_activation(x) - - for block in self.res_blocks: - x = block(x) - - if self.temporal_upsample: - x = self.upsampler(x) - # remove the first frame after upsampling. - # This is done because the first frame encodes one pixel frame. - x = x[:, :, 1:, :, :] - elif isinstance(self.upsampler, SpatialRationalResampler): - x = self.upsampler(x) - else: - x = rearrange(x, "b c f h w -> (b f) c h w") - x = self.upsampler(x) - x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) - - for block in self.post_upsample_res_blocks: - x = block(x) - - x = self.final_conv(x) - - return x - - -def upsample_video(latent: torch.Tensor, video_encoder: VideoEncoder, upsampler: "LatentUpsampler") -> torch.Tensor: - """ - Apply upsampling to the latent representation using the provided upsampler, - with normalization and un-normalization based on the video encoder's per-channel statistics. - Args: - latent: Input latent tensor of shape [B, C, F, H, W]. - video_encoder: VideoEncoder with per_channel_statistics for normalization. - upsampler: LatentUpsampler module to perform upsampling. - Returns: - torch.Tensor: Upsampled and re-normalized latent tensor. - """ - latent = video_encoder.per_channel_statistics.un_normalize(latent) - latent = upsampler(latent) - latent = video_encoder.per_channel_statistics.normalize(latent) - return latent diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model_configurator.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model_configurator.py deleted file mode 100644 index 2a714a417..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model_configurator.py +++ /dev/null @@ -1,30 +0,0 @@ -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.model.upsampler.model import LatentUpsampler - - -class LatentUpsamplerConfigurator(ModelConfigurator[LatentUpsampler]): - """ - Configurator for LatentUpsampler model. - Used to create a LatentUpsampler model from a configuration dictionary. - """ - - @classmethod - def from_config(cls: type[LatentUpsampler], config: dict) -> LatentUpsampler: - in_channels = config.get("in_channels", 128) - mid_channels = config.get("mid_channels", 512) - num_blocks_per_stage = config.get("num_blocks_per_stage", 4) - dims = config.get("dims", 3) - spatial_upsample = config.get("spatial_upsample", True) - temporal_upsample = config.get("temporal_upsample", False) - spatial_scale = config.get("spatial_scale", 2.0) - rational_resampler = config.get("rational_resampler", False) - return LatentUpsampler( - in_channels=in_channels, - mid_channels=mid_channels, - num_blocks_per_stage=num_blocks_per_stage, - dims=dims, - spatial_upsample=spatial_upsample, - temporal_upsample=temporal_upsample, - spatial_scale=spatial_scale, - rational_resampler=rational_resampler, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/pixel_shuffle.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/pixel_shuffle.py deleted file mode 100644 index 3c78f3bb4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/pixel_shuffle.py +++ /dev/null @@ -1,54 +0,0 @@ -import torch -from einops import rearrange - - -class PixelShuffleND(torch.nn.Module): - """ - N-dimensional pixel shuffle operation for upsampling tensors. - Args: - dims (int): Number of dimensions to apply pixel shuffle to. - - 1: Temporal (e.g., frames) - - 2: Spatial (e.g., height and width) - - 3: Spatiotemporal (e.g., depth, height, width) - upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension. - For dims=1, only the first value is used. - For dims=2, the first two values are used. - For dims=3, all three values are used. - The input tensor is rearranged so that the channel dimension is split into - smaller channels and upscaling factors, and the upscaling factors are moved - into the corresponding spatial/temporal dimensions. - Note: - This operation is equivalent to the patchifier operation in for the models. Consider - using this class instead. - """ - - def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)): - super().__init__() - assert dims in [1, 2, 3], "dims must be 1, 2, or 3" - self.dims = dims - self.upscale_factors = upscale_factors - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.dims == 3: - return rearrange( - x, - "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", - p1=self.upscale_factors[0], - p2=self.upscale_factors[1], - p3=self.upscale_factors[2], - ) - elif self.dims == 2: - return rearrange( - x, - "b (c p1 p2) h w -> b c (h p1) (w p2)", - p1=self.upscale_factors[0], - p2=self.upscale_factors[1], - ) - elif self.dims == 1: - return rearrange( - x, - "b (c p1) f h w -> b c (f p1) h w", - p1=self.upscale_factors[0], - ) - else: - raise ValueError(f"Unsupported dims: {self.dims}") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/res_block.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/res_block.py deleted file mode 100644 index 8b7603527..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/res_block.py +++ /dev/null @@ -1,37 +0,0 @@ -from typing import Optional - -import torch - - -class ResBlock(torch.nn.Module): - """ - Residual block with two convolutional layers, group normalization, and SiLU activation. - Args: - channels (int): Number of input and output channels. - mid_channels (Optional[int]): Number of channels in the intermediate convolution layer. Defaults to `channels` - if not specified. - dims (int): Dimensionality of the convolution (2 for Conv2d, 3 for Conv3d). Defaults to 3. - """ - - def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3): - super().__init__() - if mid_channels is None: - mid_channels = channels - - conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d - - self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1) - self.norm1 = torch.nn.GroupNorm(32, mid_channels) - self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1) - self.norm2 = torch.nn.GroupNorm(32, channels) - self.activation = torch.nn.SiLU() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.conv1(x) - x = self.norm1(x) - x = self.activation(x) - x = self.conv2(x) - x = self.norm2(x) - x = self.activation(x + residual) - return x diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/spatial_rational_resampler.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/spatial_rational_resampler.py deleted file mode 100644 index 0d2a286d3..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/spatial_rational_resampler.py +++ /dev/null @@ -1,46 +0,0 @@ -from typing import Tuple - -import torch -from einops import rearrange -from ltx_core.model.upsampler.blur_downsample import BlurDownsample -from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND - - -def _rational_for_scale(scale: float) -> Tuple[int, int]: - mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)} - if float(scale) not in mapping: - raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}") - return mapping[float(scale)] - - -class SpatialRationalResampler(torch.nn.Module): - """ - Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then anti-aliased - downsample by 'den' using fixed blur + stride. Operates on H,W only. - For dims==3, work per-frame for spatial scaling (temporal axis untouched). - Args: - mid_channels (`int`): Number of intermediate channels for the convolution layer - scale (`float`): Spatial scaling factor. Supported values are: - - 0.75: Downsample by 3/4 (reduce spatial size) - - 1.5: Upsample by 3/2 (increase spatial size) - - 2.0: Upsample by 2x (double spatial size) - - 4.0: Upsample by 4x (quadruple spatial size) - Any other value will raise a ValueError. - """ - - def __init__(self, mid_channels: int, scale: float): - super().__init__() - self.scale = float(scale) - self.num, self.den = _rational_for_scale(self.scale) - self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1) - self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) - self.blur_down = BlurDownsample(dims=2, stride=self.den) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - b, _, f, _, _ = x.shape - x = rearrange(x, "b c f h w -> (b f) c h w") - x = self.conv(x) - x = self.pixel_shuffle(x) - x = self.blur_down(x) - x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) - return x diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/__init__.py deleted file mode 100644 index 45af0259f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Video VAE package.""" - -from ltx_core.model.video_vae.memory_efficient_decode import MEMORY_EFFICIENT_DECODE -from ltx_core.model.video_vae.model_configurator import ( - VAE_DECODER_COMFY_KEYS_FILTER, - VAE_ENCODER_COMFY_KEYS_FILTER, - VideoDecoderConfigurator, - VideoEncoderConfigurator, -) -from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig, TilingConfig -from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, get_video_chunks_number - -__all__ = [ - "MEMORY_EFFICIENT_DECODE", - "VAE_DECODER_COMFY_KEYS_FILTER", - "VAE_ENCODER_COMFY_KEYS_FILTER", - "SpatialTilingConfig", - "TemporalTilingConfig", - "TilingConfig", - "VideoDecoder", - "VideoDecoderConfigurator", - "VideoEncoder", - "VideoEncoderConfigurator", - "get_video_chunks_number", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/convolution.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/convolution.py deleted file mode 100644 index ed2c057a1..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/convolution.py +++ /dev/null @@ -1,316 +0,0 @@ -from typing import Tuple, Union - -import torch -from einops import rearrange -from ltx_core.model.video_vae.enums import PaddingModeType -from torch import nn -from torch.nn import functional as F - - -def make_conv_nd( # noqa: PLR0913 - dims: Union[int, Tuple[int, int]], - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - causal: bool = False, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS, -) -> nn.Module: - if not (spatial_padding_mode == temporal_padding_mode or causal): - raise NotImplementedError("spatial and temporal padding modes must be equal") - if dims == 2: - return nn.Conv2d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - padding_mode=spatial_padding_mode.value, - ) - elif dims == 3: - if causal: - return CausalConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - dilation=dilation, - groups=groups, - bias=bias, - spatial_padding_mode=spatial_padding_mode, - ) - return nn.Conv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - padding_mode=spatial_padding_mode.value, - ) - elif dims == (2, 1): - return DualConv3d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - bias=bias, - padding_mode=spatial_padding_mode.value, - ) - else: - raise ValueError(f"unsupported dimensions: {dims}") - - -def make_linear_nd( - dims: int, - in_channels: int, - out_channels: int, - bias: bool = True, -) -> nn.Module: - if dims == 2: - return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) - elif dims in (3, (2, 1)): - return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) - else: - raise ValueError(f"unsupported dimensions: {dims}") - - -class DualConv3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: Union[int, Tuple[int, int, int]] = 1, - padding: Union[int, Tuple[int, int, int]] = 0, - dilation: Union[int, Tuple[int, int, int]] = 1, - groups: int = 1, - bias: bool = True, - padding_mode: str = "zeros", - ) -> None: - super(DualConv3d, self).__init__() - - self.in_channels = in_channels - self.out_channels = out_channels - self.padding_mode = padding_mode - # Ensure kernel_size, stride, padding, and dilation are tuples of length 3 - if isinstance(kernel_size, int): - kernel_size = (kernel_size, kernel_size, kernel_size) - if kernel_size == (1, 1, 1): - raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.") - if isinstance(stride, int): - stride = (stride, stride, stride) - if isinstance(padding, int): - padding = (padding, padding, padding) - if isinstance(dilation, int): - dilation = (dilation, dilation, dilation) - - # Set parameters for convolutions - self.groups = groups - self.bias = bias - - # Define the size of the channels after the first convolution - intermediate_channels = out_channels if in_channels < out_channels else in_channels - - # Define parameters for the first convolution - self.weight1 = nn.Parameter( - torch.Tensor( - intermediate_channels, - in_channels // groups, - 1, - kernel_size[1], - kernel_size[2], - ) - ) - self.stride1 = (1, stride[1], stride[2]) - self.padding1 = (0, padding[1], padding[2]) - self.dilation1 = (1, dilation[1], dilation[2]) - if bias: - self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels)) - else: - self.register_parameter("bias1", None) - - # Define parameters for the second convolution - self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1)) - self.stride2 = (stride[0], 1, 1) - self.padding2 = (padding[0], 0, 0) - self.dilation2 = (dilation[0], 1, 1) - if bias: - self.bias2 = nn.Parameter(torch.Tensor(out_channels)) - else: - self.register_parameter("bias2", None) - - # Initialize weights and biases - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5)) - nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5)) - if self.bias: - fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1) - bound1 = 1 / torch.sqrt(fan_in1) - nn.init.uniform_(self.bias1, -bound1, bound1) - fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2) - bound2 = 1 / torch.sqrt(fan_in2) - nn.init.uniform_(self.bias2, -bound2, bound2) - - def forward( - self, - x: torch.Tensor, - use_conv3d: bool = False, - skip_time_conv: bool = False, - ) -> torch.Tensor: - if use_conv3d: - return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv) - else: - return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv) - - def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: - # First convolution - x = F.conv3d( - x, - self.weight1, - self.bias1, - self.stride1, - self.padding1, - self.dilation1, - self.groups, - padding_mode=self.padding_mode, - ) - - if skip_time_conv: - return x - - # Second convolution - x = F.conv3d( - x, - self.weight2, - self.bias2, - self.stride2, - self.padding2, - self.dilation2, - self.groups, - padding_mode=self.padding_mode, - ) - - return x - - def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: - b, _, _, h, w = x.shape - - # First 2D convolution - x = rearrange(x, "b c d h w -> (b d) c h w") - # Squeeze the depth dimension out of weight1 since it's 1 - weight1 = self.weight1.squeeze(2) - # Select stride, padding, and dilation for the 2D convolution - stride1 = (self.stride1[1], self.stride1[2]) - padding1 = (self.padding1[1], self.padding1[2]) - dilation1 = (self.dilation1[1], self.dilation1[2]) - x = F.conv2d( - x, - weight1, - self.bias1, - stride1, - padding1, - dilation1, - self.groups, - padding_mode=self.padding_mode, - ) - - _, _, h, w = x.shape - - if skip_time_conv: - x = rearrange(x, "(b d) c h w -> b c d h w", b=b) - return x - - # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension - x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b) - - # Reshape weight2 to match the expected dimensions for conv1d - weight2 = self.weight2.squeeze(-1).squeeze(-1) - # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution - stride2 = self.stride2[0] - padding2 = self.padding2[0] - dilation2 = self.dilation2[0] - x = F.conv1d( - x, - weight2, - self.bias2, - stride2, - padding2, - dilation2, - self.groups, - padding_mode=self.padding_mode, - ) - x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w) - - return x - - @property - def weight(self) -> torch.Tensor: - return self.weight2 - - -class CausalConv3d(nn.Module): - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 3, - stride: Union[int, Tuple[int]] = 1, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ) -> None: - super().__init__() - - self.in_channels = in_channels - self.out_channels = out_channels - - kernel_size = (kernel_size, kernel_size, kernel_size) - self.time_kernel_size = kernel_size[0] - - dilation = (dilation, 1, 1) - - height_pad = kernel_size[1] // 2 - width_pad = kernel_size[2] // 2 - padding = (0, height_pad, width_pad) - - self.conv = nn.Conv3d( - in_channels, - out_channels, - kernel_size, - stride=stride, - dilation=dilation, - padding=padding, - padding_mode=spatial_padding_mode.value, - groups=groups, - bias=bias, - ) - - def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: - if causal: - first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) - x = torch.concatenate((first_frame_pad, x), dim=2) - else: - first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) - last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) - x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2) - x = self.conv(x) - return x - - @property - def weight(self) -> torch.Tensor: - return self.conv.weight diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/enums.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/enums.py deleted file mode 100644 index edb2eace1..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/enums.py +++ /dev/null @@ -1,20 +0,0 @@ -from enum import Enum - - -class NormLayerType(Enum): - GROUP_NORM = "group_norm" - PIXEL_NORM = "pixel_norm" - - -class LogVarianceType(Enum): - PER_CHANNEL = "per_channel" - UNIFORM = "uniform" - CONSTANT = "constant" - NONE = "none" - - -class PaddingModeType(Enum): - ZEROS = "zeros" - REFLECT = "reflect" - REPLICATE = "replicate" - CIRCULAR = "circular" diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/memory_efficient_decode.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/memory_efficient_decode.py deleted file mode 100644 index c55cc564b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/memory_efficient_decode.py +++ /dev/null @@ -1,649 +0,0 @@ -"""Memory-efficient VAE decoder operations. -Reduces peak VRAM usage during video decoding through in-place operations -and workspace buffer reuse. The main optimizations are: -1. **Workspace buffers** -- Pre-allocated tensors with temporal padding replace - dynamic padding (``F.pad`` / ``concatenate``) in ``CausalConv3d``. A - workspace of shape ``[B, C, T+2, H, W]`` holds the data in positions - ``[1:-1]`` with replicate padding at ``[0]`` and ``[-1]``. -2. **In-place temporal-chunked Conv3d** *(non-causal only)* -- The convolution - output is written back into the workspace buffer, avoiding a separate - output allocation. Temporal chunking with boundary save/restore ensures - correct reads despite in-place writes. -3. **In-place normalization and affine transforms** -- PixelNorm, scale/shift, - and SiLU are applied in-place on workspace views. -4. **Free-before-conv** -- For ``DepthToSpaceUpsample`` blocks the input - tensor is freed before the convolution runs so that peak VRAM never holds - input *and* output simultaneously. -Both causal and non-causal modes are supported. Non-causal mode benefits -from all four optimizations. Causal mode benefits from optimizations 1, 3, -and 4; in-place conv (2) is skipped because the asymmetric causal padding -layout prevents clean in-place overwrites. -Usage via the ``ModuleOps`` pattern (preferred):: - from ltx_core.model.video_vae import MEMORY_EFFICIENT_DECODE - builder = decoder_builder.with_module_ops( - (*decoder_builder.module_ops, MEMORY_EFFICIENT_DECODE) - ) -Or applied directly to an existing decoder:: - from ltx_core.model.video_vae.memory_efficient_decode import ( - enable_memory_efficient_decode, - ) - enable_memory_efficient_decode(decoder) -""" - -from __future__ import annotations - -import math -from typing import TYPE_CHECKING - -import torch -from einops import rearrange -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.model.common.normalization import PixelNorm -from ltx_core.model.video_vae.convolution import CausalConv3d -from ltx_core.model.video_vae.ops import unpatchify -from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D -from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample -from torch import nn -from torch.nn import functional as F - -if TYPE_CHECKING: - from ltx_core.model.video_vae.video_vae import VideoDecoder - - -# --------------------------------------------------------------------------- -# Low-level helpers -# --------------------------------------------------------------------------- - - -def _memory_format_of(t: torch.Tensor, prefer_channels_last_3d: bool = False) -> torch.memory_format: - """Pick the memory format for a workspace allocation. - When ``prefer_channels_last_3d`` is True and ``t`` is 5D, return - ``channels_last_3d`` regardless of ``t``'s current strides -- the - workspace's ``.copy_(t)`` will transcribe the data into the new layout. - This is needed because intermediate tensors inside the decoder (after - ``rearrange`` + slice + residual add in ``_upsample_forward_efficient``) - are not NHWC-contiguous, so an auto-detect helper would silently fall - back to NCHW for every workspace after the first upsample. - Otherwise fall back to inspecting ``t``: ``channels_last_3d`` if ``t`` - already uses it, else contiguous. - """ - if prefer_channels_last_3d and t.dim() == 5: - return torch.channels_last_3d - if t.dim() == 5 and t.is_contiguous(memory_format=torch.channels_last_3d): - return torch.channels_last_3d - return torch.contiguous_format - - -def _find_temporal_split_size(num_frames: int) -> int: - """Find chunk size for in-place temporal convolution. - The chunk size ensures the last chunk has at least 3 frames - (the temporal kernel size), avoiding degenerate chunks. - """ - for s in range(16, 2, -1): - remainder = num_frames % s - if remainder == 0 or remainder >= 3: - return s - - raise ValueError( - f"Unable to find a valid temporal split size for num_frames={num_frames}. Expected a split size between 3 and 16 such that the final chunk is either exact or has at least 3 frames." - ) - - -def _pad_workspace_temporal(workspace: torch.Tensor) -> None: - """Apply non-causal replicate padding to temporal boundaries. - Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and - ``workspace[:, :, -1]`` to a copy of ``workspace[:, :, -2]``. - """ - workspace[:, :, 0, :, :].copy_(workspace[:, :, 1, :, :]) - workspace[:, :, -1, :, :].copy_(workspace[:, :, -2, :, :]) - - -# --------------------------------------------------------------------------- -# In-place Conv3d (non-causal only) -# --------------------------------------------------------------------------- - - -def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) -> None: - """Run a 3x3x3 Conv3d in-place on a temporally-padded workspace. - The workspace has shape ``[B, C, T+2, H, W]`` where positions ``[1:-1]`` - hold the real data and positions ``[0]`` and ``[-1]`` are padding slots. - The convolution must have ``kernel_size=(3,3,3)``, ``stride=(1,1,1)``, - ``padding=(0,1,1)`` -- no temporal padding, symmetric spatial padding. - The output (T frames) overwrites positions ``[1:-1]``. Temporal chunking - with boundary save/restore ensures each chunk reads unmodified input even - though earlier chunks already wrote to the same buffer. - Only valid for **non-causal** mode (symmetric replicate padding). - Args: - workspace: Tensor ``[B, max(C_in, C_out), T+2, H, W]``. - Modified in-place; after the call ``workspace[:, :C_out, 1:-1]`` - holds the convolution result. - conv: ``nn.Conv3d`` with the constraints above. - """ - if conv.kernel_size != (3, 3, 3): - raise ValueError(f"Expected kernel_size=(3,3,3), got {conv.kernel_size}") - if conv.stride != (1, 1, 1): - raise ValueError(f"Expected stride=(1,1,1), got {conv.stride}") - if conv.padding != (0, 1, 1): - raise ValueError(f"Expected padding=(0,1,1), got {conv.padding}") - - _pad_workspace_temporal(workspace) - - total_frames = workspace.shape[2] - out_channels = conv.out_channels - in_channels = conv.in_channels - - if total_frames > 16: - split_size = _find_temporal_split_size(total_frames) - num_splits = (total_frames + split_size - 1) // split_size - else: - split_size = total_frames - 1 - num_splits = 1 - - # 1-frame buffers for saving / restoring boundary frames across chunks. - x_buf = torch.empty( - workspace.shape[0], - workspace.shape[1], - 1, - workspace.shape[3], - workspace.shape[4], - device=workspace.device, - dtype=workspace.dtype, - memory_format=_memory_format_of(workspace), - ) - o_buf = torch.empty_like(x_buf) - - # Helper: extract a chunk and make it contiguous. Workspace views can - # inherit strides > 2^31 from the full buffer, which makes Conv3d's - # reflect-padding path (F.pad) crash with "input tensor must fit into - # 32-bit index math". A small .clone() per chunk avoids this. - needs_clone = workspace.untyped_storage().nbytes() > (2**31 - 1) * workspace.element_size() - - def _chunk(t_start: int, t_end: int) -> torch.Tensor: - s = workspace[:, :in_channels, t_start:t_end] - return s.clone() if needs_clone else s - - # --- First chunk --- - if num_splits > 1: - # Save the boundary now so the loop below can restore it. Skipped - # when there is only one chunk: the loop never runs, and the save - # would be a wasted full HW slice copy. - x_buf[:, :, 0] = workspace[:, :, split_size - 1].clone() - workspace[:, :out_channels, 1:split_size] = conv(_chunk(0, split_size + 1)) - - # --- Remaining chunks --- - for i in range(1, num_splits): - start = i * split_size - end = min((i + 1) * split_size, total_frames - 1) - - # Save the value at start-1 (now holds previous chunk's output). - o_buf[:, :, 0] = workspace[:, :, start - 1].clone() - # Restore the original input value needed by this chunk's conv. - workspace[:, :, start - 1] = x_buf[:, :, 0] - # Save the boundary for the *next* chunk before we overwrite it. - x_buf[:, :, 0] = workspace[:, :, end - 1].clone() - - workspace[:, :out_channels, start:end] = conv(_chunk(start - 1, end + 1)) - - # Put back the previous chunk's output at the boundary. - workspace[:, :, start - 1] = o_buf[:, :, 0] - - -# --------------------------------------------------------------------------- -# Causal conv helper (free-before-conv) -# --------------------------------------------------------------------------- - - -def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor: - """Build a causal-padded buffer of shape ``[B, C, T+pad_size, H, W]``. - Copies ``x`` into ``padded[:, :, pad_size:]`` and replicates the first - real frame into the leading ``pad_size`` slots. The caller still owns - ``x`` after this returns. - """ - padded = torch.empty( - x.shape[0], - x.shape[1], - x.shape[2] + pad_size, - x.shape[3], - x.shape[4], - device=x.device, - dtype=x.dtype, - memory_format=_memory_format_of(x), - ) - padded[:, :, pad_size:].copy_(x) - for i in range(pad_size): - padded[:, :, i] = padded[:, :, pad_size] - return padded - - -def _causal_pad_free_and_conv(x: torch.Tensor, causal_conv: CausalConv3d) -> torch.Tensor: - """Causal-pad *x*, free it, then run the raw ``nn.Conv3d``. - This avoids the peak where both the original and padded tensors are - live simultaneously (as happens inside ``CausalConv3d.forward``). - Args: - x: Input ``[B, C_in, T, H, W]``. **Deleted** inside this function; - the caller must not use it afterwards. - Returns: - Convolution output ``[B, C_out, T, H, W]``. - """ - padded = _causal_pad(x, causal_conv.time_kernel_size - 1) - del x - result = causal_conv.conv(padded) - del padded - return result - - -# --------------------------------------------------------------------------- -# In-place normalization -# --------------------------------------------------------------------------- - - -def _pixel_norm_inplace(x: torch.Tensor, eps: float = 1e-8) -> None: - """In-place RMS (pixel) normalization along the channel dimension.""" - rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps) - x.div_(rms) - - -def _norm_inplace(norm: nn.Module, x: torch.Tensor) -> None: - """Apply *norm* in-place, using an optimised path for ``PixelNorm``.""" - if isinstance(norm, PixelNorm): - _pixel_norm_inplace(x, eps=norm.eps) - else: - # GroupNorm or other -- fall back to allocating a temporary. - result = norm(x) - x.copy_(result) - del result - - -# --------------------------------------------------------------------------- -# Per-block efficient forwards -# --------------------------------------------------------------------------- - - -def _resnet_block_forward_inplace( - resnet: ResnetBlock3D, - workspace: torch.Tensor, - causal: bool, - timestep: torch.Tensor | None, - generator: torch.Generator | None, -) -> None: - """Run a ``ResnetBlock3D`` in-place on a workspace buffer. - The workspace has shape ``[B, C, T+2, H, W]`` with real data in - ``[1:-1]``. After this call ``workspace[:, :, 1:-1]`` holds the - residual-branch output ``F(x)`` (without the skip connection -- - the caller adds it back to the hidden state). - Only valid when ``in_channels == out_channels`` (true for all - ``ResnetBlock3D`` instances inside a ``UNetMidBlock3D``). - """ - if resnet.in_channels != resnet.out_channels: - raise ValueError(f"In-place resnet forward requires in_channels == out_channels, got {resnet.in_channels} != {resnet.out_channels}") - - interior = workspace[:, :, 1:-1] - - # --- norm1 + [ada scaling] + SiLU + conv1 --- - _norm_inplace(resnet.norm1, interior) - - if resnet.timestep_conditioning and timestep is not None: - ada = resnet.scale_shift_table[None, ..., None, None, None].to(device=interior.device, dtype=interior.dtype) + timestep.reshape( - interior.shape[0], - 4, - -1, - timestep.shape[-3], - timestep.shape[-2], - timestep.shape[-1], - ) - shift1, scale1, shift2, scale2 = ada.unbind(dim=1) - interior.mul_(1 + scale1).add_(shift1) - - F.silu(interior, inplace=True) - - if causal: - result = resnet.conv1(interior, causal=True) - interior.copy_(result) - del result - else: - inplace_conv3d_temporal_chunked(workspace, resnet.conv1.conv) - - if resnet.inject_noise: - spatial_shape = interior.shape[-2:] - scale = resnet.per_channel_scale1.to(device=interior.device, dtype=interior.dtype) - noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator) - interior.add_((noise * scale)[None, :, None, ...]) - - # --- norm2 + [ada scaling] + SiLU + conv2 --- - _norm_inplace(resnet.norm2, interior) - - if resnet.timestep_conditioning and timestep is not None: - interior.mul_(1 + scale2).add_(shift2) # type: ignore[possibly-undefined] - - F.silu(interior, inplace=True) - # dropout is always 0.0 during inference -- skip. - - if causal: - result = resnet.conv2(interior, causal=True) - interior.copy_(result) - del result - else: - inplace_conv3d_temporal_chunked(workspace, resnet.conv2.conv) - - if resnet.inject_noise: - spatial_shape = interior.shape[-2:] - scale = resnet.per_channel_scale2.to(device=interior.device, dtype=interior.dtype) - noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator) - interior.add_((noise * scale)[None, :, None, ...]) - - -def _midblock_forward_efficient( - block: UNetMidBlock3D, - hidden_states: torch.Tensor, - causal: bool, - timestep: torch.Tensor | None, - generator: torch.Generator | None, - prefer_channels_last_3d: bool = False, -) -> torch.Tensor: - """Memory-efficient ``UNetMidBlock3D`` forward. - Allocates a single workspace buffer that is reused across all - ``ResnetBlock3D`` iterations. For each block the workspace is - populated with the current hidden state, processed in-place, and - the result is added back (residual connection). - """ - timestep_embed = None - if block.timestep_conditioning: - if timestep is None: - raise ValueError("'timestep' required when timestep_conditioning=True") - batch_size = hidden_states.shape[0] - timestep_embed = block.time_embedder( - timestep=timestep.flatten(), - hidden_dtype=hidden_states.dtype, - ) - timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) - - workspace = torch.empty( - hidden_states.shape[0], - hidden_states.shape[1], - hidden_states.shape[2] + 2, - hidden_states.shape[3], - hidden_states.shape[4], - device=hidden_states.device, - dtype=hidden_states.dtype, - memory_format=_memory_format_of(hidden_states, prefer_channels_last_3d), - ) - - for resnet in block.res_blocks: - workspace[:, :, 1:-1].copy_(hidden_states) - _resnet_block_forward_inplace(resnet, workspace, causal, timestep_embed, generator) - hidden_states.add_(workspace[:, :, 1:-1]) - - del workspace - return hidden_states - - -def _upsample_forward_efficient( - block: DepthToSpaceUpsample, - x: torch.Tensor, - causal: bool, - prefer_channels_last_3d: bool = False, -) -> torch.Tensor: - """Memory-efficient ``DepthToSpaceUpsample`` forward. - For non-causal mode the input is copied into a workspace and the - convolution runs in-place. For causal mode the input is manually - padded and freed before the convolution runs. Both paths avoid - the peak where input *and* output coexist. - """ - if block.residual: - x_in = rearrange( - x, - "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", - p1=block.stride[0], - p2=block.stride[1], - p3=block.stride[2], - ) - num_repeat = math.prod(block.stride) // block.out_channels_reduction_factor - x_in = x_in.repeat(1, num_repeat, 1, 1, 1) - if block.stride[0] == 2: - x_in = x_in[:, :, 1:, :, :] - - conv = block.conv.conv # underlying nn.Conv3d inside CausalConv3d - in_channels = x.shape[1] - out_channels = conv.out_channels - - if causal: - x = _causal_pad_free_and_conv(x, block.conv) - else: - mem_fmt = _memory_format_of(x, prefer_channels_last_3d) - workspace = torch.empty( - x.shape[0], - max(in_channels, out_channels), - x.shape[2] + 2, - x.shape[3], - x.shape[4], - device=x.device, - dtype=x.dtype, - memory_format=mem_fmt, - ) - workspace[:, :in_channels, 1:-1].copy_(x) - del x - inplace_conv3d_temporal_chunked(workspace, conv) - x = workspace[:, :out_channels, 1:-1].contiguous(memory_format=mem_fmt) - del workspace - - x = rearrange( - x, - "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", - p1=block.stride[0], - p2=block.stride[1], - p3=block.stride[2], - ) - if block.stride[0] == 2: - x = x[:, :, 1:, :, :] - if block.residual: - x.add_(x_in) - del x_in - return x - - -# --------------------------------------------------------------------------- -# Final norm + conv_out -# --------------------------------------------------------------------------- - - -def _final_norm_and_conv_out( - decoder: VideoDecoder, - sample: torch.Tensor, - causal: bool, - scaled_timestep: torch.Tensor | None, - batch_size: int, - prefer_channels_last_3d: bool = False, -) -> torch.Tensor: - """Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify.""" - conv_out_mod: CausalConv3d = decoder.conv_out # type: ignore[assignment] - conv_out = conv_out_mod.conv - feature_channels = sample.shape[1] - - mem_fmt = _memory_format_of(sample, prefer_channels_last_3d) - workspace = torch.empty( - sample.shape[0], - max(feature_channels, conv_out.out_channels), - sample.shape[2] + 2, - sample.shape[3], - sample.shape[4], - device=sample.device, - dtype=sample.dtype, - memory_format=mem_fmt, - ) - workspace[:, :feature_channels, 1:-1].copy_(sample) - del sample - - interior = workspace[:, :feature_channels, 1:-1] - _norm_inplace(decoder.conv_norm_out, interior) - - if decoder.timestep_conditioning: - embedded_timestep = decoder.last_time_embedder( - timestep=scaled_timestep.flatten(), - hidden_dtype=interior.dtype, - ) - embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) - ada_values = decoder.last_scale_shift_table[None, ..., None, None, None].to(device=interior.device, dtype=interior.dtype) + embedded_timestep.reshape( - batch_size, - 2, - -1, - embedded_timestep.shape[-3], - embedded_timestep.shape[-2], - embedded_timestep.shape[-1], - ) - shift, scale = ada_values.unbind(dim=1) - interior.mul_(1 + scale).add_(shift) - - F.silu(interior, inplace=True) - - if causal: - # Causal: build padded tensor directly from the interior view, - # then free the workspace before running the conv. - padded = _causal_pad(interior, conv_out_mod.time_kernel_size - 1) - del workspace, interior - result = conv_out(padded) - del padded - else: - inplace_conv3d_temporal_chunked(workspace, conv_out) - result = workspace[:, : conv_out.out_channels, 1:-1].contiguous(memory_format=mem_fmt) - del workspace, interior - - return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1) - - -# --------------------------------------------------------------------------- -# Top-level efficient decoder forward -# --------------------------------------------------------------------------- - - -def _memory_efficient_forward( - decoder: VideoDecoder, - sample: torch.Tensor, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, -) -> torch.Tensor: - """Full memory-efficient ``VideoDecoder.forward`` replacement. - Orchestrates the entire decode through workspace-based operations: - ``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient - paths; standalone ``ResnetBlock3D`` blocks fall back to the standard - forward. The final norm + ada + SiLU + conv_out is also workspace-based. - All workspaces are allocated ``channels_last_3d`` so cuDNN's NHWC 3D - conv kernels run end-to-end. The caller (:func:`enable_memory_efficient_decode`) - is responsible for converting the input sample and decoder weights to NHWC. - """ - causal = decoder.causal - batch_size = sample.shape[0] - sample = sample.to(next(decoder.parameters()).dtype) - - # --- Noise injection and de-normalisation (identical to standard path) --- - if decoder.timestep_conditioning: - noise = torch.randn(sample.size(), generator=generator, dtype=sample.dtype, device=sample.device) * decoder.decode_noise_scale - sample = noise + (1.0 - decoder.decode_noise_scale) * sample - - sample = decoder.per_channel_statistics.un_normalize(sample) - - if timestep is None and decoder.timestep_conditioning: - timestep = torch.full((batch_size,), decoder.decode_timestep, device=sample.device, dtype=sample.dtype) - - # --- conv_in (latent tensor is small -- standard path is fine) --- - sample = decoder.conv_in(sample, causal=causal) - - upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype - sample = sample.to(upscale_dtype) - - scaled_timestep = None - if decoder.timestep_conditioning: - if timestep is None: - raise ValueError("'timestep' required when timestep_conditioning=True") - scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample) - - # Workspaces are unconditionally NHWC: rearrange + slice + residual-add - # inside _upsample_forward_efficient produces NCHW-default output, so - # per-tensor inspection would silently fall back to NCHW for every - # workspace after the first upsample. - - # --- Up blocks (dispatch to efficient path per block type) --- - for up_block in decoder.up_blocks: - if isinstance(up_block, UNetMidBlock3D): - sample = _midblock_forward_efficient( - up_block, - sample, - causal=causal, - timestep=scaled_timestep if decoder.timestep_conditioning else None, - generator=generator, - prefer_channels_last_3d=True, - ) - elif isinstance(up_block, DepthToSpaceUpsample): - sample = _upsample_forward_efficient(up_block, sample, causal=causal, prefer_channels_last_3d=True) - elif isinstance(up_block, ResnetBlock3D): - sample = up_block(sample, causal=causal, generator=generator) - else: - sample = up_block(sample, causal=causal) - - return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size, prefer_channels_last_3d=True) - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - - -def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module: - """Patch a ``VideoDecoder`` to use the memory-efficient forward path. - The mem-efficient path runs the decoder in ``channels_last_3d`` memory - format: weights and inputs are converted on first call so cuDNN's NHWC - 3D conv kernels are used (~2x faster, avoids the large vol2col scratch - buffer of the NCHW path). - The original ``forward`` is saved as ``decoder._original_forward`` so - that it can be restored later with :func:`disable_memory_efficient_decode`. - """ - # Import here to avoid circular dependency at module level. - from ltx_core.model.video_vae.video_vae import VideoDecoder # noqa: PLC0415 - - if not isinstance(decoder, VideoDecoder): - raise TypeError(f"Expected VideoDecoder, got {type(decoder).__name__}") - - if hasattr(decoder, "_original_forward"): - return decoder - - original_forward = decoder.forward - weights_converted = False - - def efficient_forward( - sample: torch.Tensor, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - nonlocal weights_converted - if sample.dim() == 5: - if not weights_converted: - # Lazy: weights are real by first-call time (meta -> loader -> here). - decoder.to(memory_format=torch.channels_last_3d) - weights_converted = True - sample = sample.to(memory_format=torch.channels_last_3d) - return _memory_efficient_forward(decoder, sample, timestep, generator) - - decoder._original_forward = original_forward # type: ignore[attr-defined] - decoder.forward = efficient_forward # type: ignore[assignment] - return decoder - - -def disable_memory_efficient_decode(decoder: nn.Module) -> nn.Module: - """Restore the original ``forward`` method on a patched ``VideoDecoder``.""" - if hasattr(decoder, "_original_forward"): - decoder.forward = decoder._original_forward # type: ignore[attr-defined] - del decoder._original_forward # type: ignore[attr-defined] - return decoder - - -def _is_video_decoder(model: nn.Module) -> bool: - """Matcher for the ``MEMORY_EFFICIENT_DECODE`` module op.""" - from ltx_core.model.video_vae.video_vae import VideoDecoder # noqa: PLC0415 - - return isinstance(model, VideoDecoder) - - -MEMORY_EFFICIENT_DECODE = ModuleOps( - name="memory_efficient_vae_decode", - matcher=_is_video_decoder, - mutator=enable_memory_efficient_decode, -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/model_configurator.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/model_configurator.py deleted file mode 100644 index 8e64ce5c8..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/model_configurator.py +++ /dev/null @@ -1,79 +0,0 @@ -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType -from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder - - -class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]): - """Configurator for creating a video VAE Encoder from a configuration dictionary.""" - - @classmethod - def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder: - config = config.get("vae", {}) - convolution_dimensions = config.get("dims", 3) - in_channels = config.get("in_channels", 3) - latent_channels = config.get("latent_channels", 128) - spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "zeros")) - encoder_blocks = config.get("encoder_blocks", []) - patch_size = config.get("patch_size", 4) - norm_layer_str = config.get("norm_layer", "pixel_norm") - latent_log_var_str = config.get("latent_log_var", "uniform") - - return VideoEncoder( - convolution_dimensions=convolution_dimensions, - in_channels=in_channels, - out_channels=latent_channels, - encoder_blocks=encoder_blocks, - patch_size=patch_size, - norm_layer=NormLayerType(norm_layer_str), - latent_log_var=LogVarianceType(latent_log_var_str), - encoder_spatial_padding_mode=spatial_padding_mode, - ) - - -class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]): - """Configurator for creating a video VAE Decoder from a configuration dictionary.""" - - @classmethod - def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder: - config = config.get("vae", {}) - convolution_dimensions = config.get("dims", 3) - latent_channels = config.get("latent_channels", 128) - spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "reflect")) - out_channels = config.get("out_channels", 3) - decoder_blocks = config.get("decoder_blocks", []) - patch_size = config.get("patch_size", 4) - norm_layer_str = config.get("norm_layer", "pixel_norm") - causal = config.get("causal_decoder", False) - timestep_conditioning = config.get("timestep_conditioning", True) - base_channels = config.get("decoder_base_channels", 128) - - return VideoDecoder( - convolution_dimensions=convolution_dimensions, - in_channels=latent_channels, - out_channels=out_channels, - decoder_blocks=decoder_blocks, - patch_size=patch_size, - norm_layer=NormLayerType(norm_layer_str), - causal=causal, - timestep_conditioning=timestep_conditioning, - decoder_spatial_padding_mode=spatial_padding_mode, - base_channels=base_channels, - ) - - -VAE_DECODER_COMFY_KEYS_FILTER = ( - SDOps("VAE_DECODER_COMFY_KEYS_FILTER") - .with_matching(prefix="vae.decoder.") - .with_matching(prefix="vae.per_channel_statistics.") - .with_replacement("vae.decoder.", "") - .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.") -) - -VAE_ENCODER_COMFY_KEYS_FILTER = ( - SDOps("VAE_ENCODER_COMFY_KEYS_FILTER") - .with_matching(prefix="vae.encoder.") - .with_matching(prefix="vae.per_channel_statistics.") - .with_replacement("vae.encoder.", "") - .with_replacement("vae.per_channel_statistics.", "per_channel_statistics.") -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/normalization.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/normalization.py deleted file mode 100644 index 45c003533..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/normalization.py +++ /dev/null @@ -1,3 +0,0 @@ -from ltx_core.model.common.normalization import PixelNorm, build_normalization_layer - -__all__ = ["PixelNorm", "build_normalization_layer"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/ops.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/ops.py deleted file mode 100644 index 022b42460..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/ops.py +++ /dev/null @@ -1,78 +0,0 @@ -import torch -from einops import rearrange -from torch import nn - - -def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - """ - Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks - and moves pixels from each block into separate channels (space-to-depth). - Args: - x: Input tensor (4D or 5D) - patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. - patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). - For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) - Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) - """ - if patch_size_hw == 1 and patch_size_t == 1: - return x - if x.dim() == 4: - x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) - elif x.dim() == 5: - x = rearrange( - x, - "b c (f p) (h q) (w r) -> b (c p r q) f h w", - p=patch_size_t, - q=patch_size_hw, - r=patch_size_hw, - ) - else: - raise ValueError(f"Invalid input shape: {x.shape}") - - return x - - -def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: - """ - Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from - channels back into patch_size x patch_size blocks (depth-to-space). - Args: - x: Input tensor (4D or 5D) - patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. - patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). - For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) - Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) - """ - if patch_size_hw == 1 and patch_size_t == 1: - return x - - if x.dim() == 4: - x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) - elif x.dim() == 5: - x = rearrange( - x, - "b (c p r q) f h w -> b c (f p) (h q) (w r)", - p=patch_size_t, - q=patch_size_hw, - r=patch_size_hw, - ) - - return x - - -class PerChannelStatistics(nn.Module): - """ - Per-channel statistics for normalizing and denormalizing the latent representation. - This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. - """ - - def __init__(self, latent_channels: int = 128): - super().__init__() - self.register_buffer("std-of-means", torch.empty(latent_channels)) - self.register_buffer("mean-of-means", torch.empty(latent_channels)) - - def un_normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x) - - def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/resnet.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/resnet.py deleted file mode 100644 index 3dd440f29..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/resnet.py +++ /dev/null @@ -1,264 +0,0 @@ -from typing import Optional, Tuple, Union - -import torch -from ltx_core.model.common.normalization import PixelNorm -from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings -from ltx_core.model.video_vae.convolution import make_conv_nd, make_linear_nd -from ltx_core.model.video_vae.enums import NormLayerType, PaddingModeType -from torch import nn - - -class ResnetBlock3D(nn.Module): - r""" - A Resnet block. - Parameters: - in_channels (`int`): The number of channels in the input. - out_channels (`int`, *optional*, default to be `None`): - The number of output channels for the first conv layer. If None, same as `in_channels`. - dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. - groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. - eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. - """ - - def __init__( - self, - dims: Union[int, Tuple[int, int]], - in_channels: int, - out_channels: Optional[int] = None, - dropout: float = 0.0, - groups: int = 32, - eps: float = 1e-6, - norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, - inject_noise: bool = False, - timestep_conditioning: bool = False, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.inject_noise = inject_noise - - if norm_layer == NormLayerType.GROUP_NORM: - self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) - elif norm_layer == NormLayerType.PIXEL_NORM: - self.norm1 = PixelNorm() - - self.non_linearity = nn.SiLU() - - self.conv1 = make_conv_nd( - dims, - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - - if inject_noise: - self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1))) - - if norm_layer == NormLayerType.GROUP_NORM: - self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) - elif norm_layer == NormLayerType.PIXEL_NORM: - self.norm2 = PixelNorm() - - self.dropout = torch.nn.Dropout(dropout) - - self.conv2 = make_conv_nd( - dims, - out_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - - if inject_noise: - self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1))) - - self.conv_shortcut = make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) if in_channels != out_channels else nn.Identity() - - # Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout - # avoiding the need for dimension rearrangement used in standard nn.LayerNorm - self.norm3 = nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) if in_channels != out_channels else nn.Identity() - - self.timestep_conditioning = timestep_conditioning - - if timestep_conditioning: - self.scale_shift_table = nn.Parameter(torch.zeros(4, in_channels)) - - def _feed_spatial_noise( - self, - hidden_states: torch.Tensor, - per_channel_scale: torch.Tensor, - generator: Optional[torch.Generator] = None, - ) -> torch.Tensor: - spatial_shape = hidden_states.shape[-2:] - device = hidden_states.device - dtype = hidden_states.dtype - - # similar to the "explicit noise inputs" method in style-gan - spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None] - scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...] - hidden_states = hidden_states + scaled_noise - - return hidden_states - - def forward( - self, - input_tensor: torch.Tensor, - causal: bool = True, - timestep: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> torch.Tensor: - hidden_states = input_tensor - batch_size = hidden_states.shape[0] - - hidden_states = self.norm1(hidden_states) - if self.timestep_conditioning: - if timestep is None: - raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") - ada_values = self.scale_shift_table[None, ..., None, None, None].to(device=hidden_states.device, dtype=hidden_states.dtype) + timestep.reshape( - batch_size, - 4, - -1, - timestep.shape[-3], - timestep.shape[-2], - timestep.shape[-1], - ) - shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) - - hidden_states = hidden_states * (1 + scale1) + shift1 - - hidden_states = self.non_linearity(hidden_states) - - hidden_states = self.conv1(hidden_states, causal=causal) - - if self.inject_noise: - hidden_states = self._feed_spatial_noise( - hidden_states, - self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype), - generator=generator, - ) - - hidden_states = self.norm2(hidden_states) - - if self.timestep_conditioning: - hidden_states = hidden_states * (1 + scale2) + shift2 - - hidden_states = self.non_linearity(hidden_states) - - hidden_states = self.dropout(hidden_states) - - hidden_states = self.conv2(hidden_states, causal=causal) - - if self.inject_noise: - hidden_states = self._feed_spatial_noise( - hidden_states, - self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype), - generator=generator, - ) - - input_tensor = self.norm3(input_tensor) - - batch_size = input_tensor.shape[0] - - input_tensor = self.conv_shortcut(input_tensor) - - output_tensor = input_tensor + hidden_states - - return output_tensor - - -class UNetMidBlock3D(nn.Module): - """ - A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks. - Args: - in_channels (`int`): The number of input channels. - dropout (`float`, *optional*, defaults to 0.0): The dropout rate. - num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. - resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. - resnet_groups (`int`, *optional*, defaults to 32): - The number of groups to use in the group normalization layers of the resnet blocks. - norm_layer (`str`, *optional*, defaults to `group_norm`): - The normalization layer to use. Can be either `group_norm` or `pixel_norm`. - inject_noise (`bool`, *optional*, defaults to `False`): - Whether to inject noise into the hidden states. - timestep_conditioning (`bool`, *optional*, defaults to `False`): - Whether to condition the hidden states on the timestep. - Returns: - `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, - in_channels, height, width)`. - """ - - def __init__( - self, - dims: Union[int, Tuple[int, int]], - in_channels: int, - dropout: float = 0.0, - num_layers: int = 1, - resnet_eps: float = 1e-6, - resnet_groups: int = 32, - norm_layer: NormLayerType = NormLayerType.GROUP_NORM, - inject_noise: bool = False, - timestep_conditioning: bool = False, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ): - super().__init__() - resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) - - self.timestep_conditioning = timestep_conditioning - - if timestep_conditioning: - self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=in_channels * 4, size_emb_dim=0) - - self.res_blocks = nn.ModuleList( - [ - ResnetBlock3D( - dims=dims, - in_channels=in_channels, - out_channels=in_channels, - eps=resnet_eps, - groups=resnet_groups, - dropout=dropout, - norm_layer=norm_layer, - inject_noise=inject_noise, - timestep_conditioning=timestep_conditioning, - spatial_padding_mode=spatial_padding_mode, - ) - for _ in range(num_layers) - ] - ) - - def forward( - self, - hidden_states: torch.Tensor, - causal: bool = True, - timestep: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - ) -> torch.Tensor: - timestep_embed = None - if self.timestep_conditioning: - if timestep is None: - raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") - batch_size = hidden_states.shape[0] - timestep_embed = self.time_embedder( - timestep=timestep.flatten(), - hidden_dtype=hidden_states.dtype, - ) - timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) - - for resnet in self.res_blocks: - hidden_states = resnet( - hidden_states, - causal=causal, - timestep=timestep_embed, - generator=generator, - ) - - return hidden_states diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/sampling.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/sampling.py deleted file mode 100644 index 5d8f7427b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/sampling.py +++ /dev/null @@ -1,123 +0,0 @@ -import math -from typing import Tuple, Union - -import torch -from einops import rearrange -from torch import nn - -from .convolution import make_conv_nd -from .enums import PaddingModeType - - -class SpaceToDepthDownsample(nn.Module): - def __init__( - self, - dims: Union[int, Tuple[int, int]], - in_channels: int, - out_channels: int, - stride: Tuple[int, int, int], - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ): - super().__init__() - self.stride = stride - self.group_size = in_channels * math.prod(stride) // out_channels - self.conv = make_conv_nd( - dims=dims, - in_channels=in_channels, - out_channels=out_channels // math.prod(stride), - kernel_size=3, - stride=1, - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - - def forward( - self, - x: torch.Tensor, - causal: bool = True, - ) -> torch.Tensor: - if self.stride[0] == 2: - x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding - - # skip connection - x_in = rearrange( - x, - "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", - p1=self.stride[0], - p2=self.stride[1], - p3=self.stride[2], - ) - x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size) - x_in = x_in.mean(dim=2) - - # conv - x = self.conv(x, causal=causal) - x = rearrange( - x, - "b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w", - p1=self.stride[0], - p2=self.stride[1], - p3=self.stride[2], - ) - - x = x + x_in - - return x - - -class DepthToSpaceUpsample(nn.Module): - def __init__( - self, - dims: int | Tuple[int, int], - in_channels: int, - stride: Tuple[int, int, int], - residual: bool = False, - out_channels_reduction_factor: int = 1, - spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ): - super().__init__() - self.stride = stride - self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor - self.conv = make_conv_nd( - dims=dims, - in_channels=in_channels, - out_channels=self.out_channels, - kernel_size=3, - stride=1, - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - self.residual = residual - self.out_channels_reduction_factor = out_channels_reduction_factor - - def forward( - self, - x: torch.Tensor, - causal: bool = True, - ) -> torch.Tensor: - if self.residual: - # Reshape and duplicate the input to match the output shape - x_in = rearrange( - x, - "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", - p1=self.stride[0], - p2=self.stride[1], - p3=self.stride[2], - ) - num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor - x_in = x_in.repeat(1, num_repeat, 1, 1, 1) - if self.stride[0] == 2: - x_in = x_in[:, :, 1:, :, :] - x = self.conv(x, causal=causal) - x = rearrange( - x, - "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", - p1=self.stride[0], - p2=self.stride[1], - p3=self.stride[2], - ) - if self.stride[0] == 2: - x = x[:, :, 1:, :, :] - if self.residual: - x = x + x_in - return x diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/tiling.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/tiling.py deleted file mode 100644 index 8019a45d9..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/tiling.py +++ /dev/null @@ -1,65 +0,0 @@ -from dataclasses import dataclass - - -@dataclass(frozen=True) -class SpatialTilingConfig: - """Configuration for dividing each frame into spatial tiles with optional overlap. - Args: - tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32. - tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0. - """ - - tile_size_in_pixels: int - tile_overlap_in_pixels: int = 0 - - def __post_init__(self) -> None: - if self.tile_size_in_pixels < 64: - raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}") - if self.tile_size_in_pixels % 32 != 0: - raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}") - if self.tile_overlap_in_pixels % 32 != 0: - raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}") - if self.tile_overlap_in_pixels >= self.tile_size_in_pixels: - raise ValueError(f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}") - - -@dataclass(frozen=True) -class TemporalTilingConfig: - """Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap. - Args: - tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8. - tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles. - Must be divisible by 8. Defaults to 0. - """ - - tile_size_in_frames: int - tile_overlap_in_frames: int = 0 - - def __post_init__(self) -> None: - if self.tile_size_in_frames < 16: - raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}") - if self.tile_size_in_frames % 8 != 0: - raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}") - if self.tile_overlap_in_frames % 8 != 0: - raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}") - if self.tile_overlap_in_frames >= self.tile_size_in_frames: - raise ValueError(f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}") - - -@dataclass(frozen=True) -class TilingConfig: - """Configuration for splitting video into tiles with optional overlap. - Attributes: - spatial_config: Configuration for splitting spatial dimensions into tiles. - temporal_config: Configuration for splitting temporal dimension into tiles. - """ - - spatial_config: SpatialTilingConfig | None = None - temporal_config: TemporalTilingConfig | None = None - - @classmethod - def default(cls) -> "TilingConfig": - return cls( - spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64), - temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24), - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/video_vae.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/video_vae.py deleted file mode 100644 index 7cef1dfd4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/video_vae.py +++ /dev/null @@ -1,1093 +0,0 @@ -import logging -from typing import Any, Callable, Iterator, List, Tuple - -import torch -from einops import rearrange -from ltx_core.model.common.normalization import PixelNorm -from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings -from ltx_core.model.video_vae.convolution import make_conv_nd -from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType -from ltx_core.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify -from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D -from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample, SpaceToDepthDownsample -from ltx_core.model.video_vae.tiling import TilingConfig -from ltx_core.tiling import ( - DEFAULT_MAPPING_OPERATION, - DEFAULT_SPLIT_OPERATION, - DimensionIntervals, - MappingOperation, - Tile, - compute_rectangular_mask_1d, - compute_trapezoidal_mask_1d, - create_tiles, - split_temporal, -) -from ltx_core.tiling import ( - split_by_size as split_in_spatial, -) -from ltx_core.tiling import ( - split_temporal_causal as split_in_temporal, -) -from ltx_core.types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors, VideoLatentShape -from torch import nn - -logger: logging.Logger = logging.getLogger(__name__) - - -def _make_encoder_block( - block_name: str, - block_config: dict[str, Any], - in_channels: int, - convolution_dimensions: int, - norm_layer: NormLayerType, - norm_num_groups: int, - spatial_padding_mode: PaddingModeType, -) -> Tuple[nn.Module, int]: - out_channels = in_channels - - if block_name == "res_x": - block = UNetMidBlock3D( - dims=convolution_dimensions, - in_channels=in_channels, - num_layers=block_config["num_layers"], - resnet_eps=1e-6, - resnet_groups=norm_num_groups, - norm_layer=norm_layer, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "res_x_y": - out_channels = in_channels * block_config.get("multiplier", 2) - block = ResnetBlock3D( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - eps=1e-6, - groups=norm_num_groups, - norm_layer=norm_layer, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_time": - block = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=(2, 1, 1), - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_space": - block = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=(1, 2, 2), - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_all": - block = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=(2, 2, 2), - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_all_x_y": - out_channels = in_channels * block_config.get("multiplier", 2) - block = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - kernel_size=3, - stride=(2, 2, 2), - causal=True, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_all_res": - out_channels = in_channels * block_config.get("multiplier", 2) - block = SpaceToDepthDownsample( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - stride=(2, 2, 2), - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_space_res": - out_channels = in_channels * block_config.get("multiplier", 2) - block = SpaceToDepthDownsample( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - stride=(1, 2, 2), - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_time_res": - out_channels = in_channels * block_config.get("multiplier", 2) - block = SpaceToDepthDownsample( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - stride=(2, 1, 1), - spatial_padding_mode=spatial_padding_mode, - ) - else: - raise ValueError(f"unknown block: {block_name}") - - return block, out_channels - - -class VideoEncoder(nn.Module): - _DEFAULT_NORM_NUM_GROUPS = 32 - """ - Variational Autoencoder Encoder. Encodes video frames into a latent representation. - The encoder compresses the input video through a series of downsampling operations controlled by - patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W'). - Compression Behavior: - The total compression is determined by: - 1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4) - 2. Sequential compression through encoder_blocks based on their stride patterns - Compression blocks apply 2x compression in specified dimensions: - - "compress_time" / "compress_time_res": temporal only - - "compress_space" / "compress_space_res": spatial only (H and W) - - "compress_all" / "compress_all_res": all dimensions (F, H, W) - - "res_x" / "res_x_y": no compression - Standard LTX Video configuration: - - patch_size=4 - - encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res - - Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32 - - Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16) - - Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...) - Args: - convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D). - in_channels: The number of input channels. For RGB images, this is 3. - out_channels: The number of output channels (latent channels). For latent channels, this is 128. - encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params) - where params is either an int (num_layers) or a dict with configuration. - patch_size: The patch size for initial spatial compression. Should be a power of 2. - norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`. - latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`. - """ - - def __init__( - self, - convolution_dimensions: int = 3, - in_channels: int = 3, - out_channels: int = 128, - encoder_blocks: List[Tuple[str, int]] | List[Tuple[str, dict[str, Any]]] = [], # noqa: B006 - patch_size: int = 4, - norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, - latent_log_var: LogVarianceType = LogVarianceType.UNIFORM, - encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, - ): - super().__init__() - - self.patch_size = patch_size - self.norm_layer = norm_layer - self.latent_channels = out_channels - self.latent_log_var = latent_log_var - self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS - - # Per-channel statistics for normalizing latents - self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels) - - in_channels = in_channels * patch_size**2 - feature_channels = out_channels - - self.conv_in = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=feature_channels, - kernel_size=3, - stride=1, - padding=1, - causal=True, - spatial_padding_mode=encoder_spatial_padding_mode, - ) - - self.down_blocks = nn.ModuleList([]) - - for block_name, block_params in encoder_blocks: - # Convert int to dict format for uniform handling - block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params - - block, feature_channels = _make_encoder_block( - block_name=block_name, - block_config=block_config, - in_channels=feature_channels, - convolution_dimensions=convolution_dimensions, - norm_layer=norm_layer, - norm_num_groups=self._norm_num_groups, - spatial_padding_mode=encoder_spatial_padding_mode, - ) - - self.down_blocks.append(block) - - # out - if norm_layer == NormLayerType.GROUP_NORM: - self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) - elif norm_layer == NormLayerType.PIXEL_NORM: - self.conv_norm_out = PixelNorm() - - self.conv_act = nn.SiLU() - - conv_out_channels = out_channels - if latent_log_var == LogVarianceType.PER_CHANNEL: - conv_out_channels *= 2 - elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}: - conv_out_channels += 1 - elif latent_log_var != LogVarianceType.NONE: - raise ValueError(f"Invalid latent_log_var: {latent_log_var}") - - self.conv_out = make_conv_nd( - dims=convolution_dimensions, - in_channels=feature_channels, - out_channels=conv_out_channels, - kernel_size=3, - padding=1, - causal=True, - spatial_padding_mode=encoder_spatial_padding_mode, - ) - - def forward(self, sample: torch.Tensor) -> torch.Tensor: - r""" - Encode video frames into normalized latent representation. - Args: - sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...). - If not, the encoder crops the last frames to the nearest valid length. - Should be normalized to [-1, 1] range before encoding. - Returns: - Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32. - Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16). - """ - # Validate frame count (crop to nearest valid length if needed) - frames_count = sample.shape[2] - if ((frames_count - 1) % 8) != 0: - frames_to_crop = (frames_count - 1) % 8 - logger.warning( - "Invalid number of frames %s for encode; cropping last %s frames to satisfy 1 + 8*k.", - frames_count, - frames_to_crop, - ) - sample = sample[:, :, :-frames_to_crop, ...] - - # Initial spatial compression: trade spatial resolution for channel depth - # This reduces H,W by patch_size and increases channels, making convolutions more efficient - # Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4 - sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) - sample = self.conv_in(sample) - - for down_block in self.down_blocks: - sample = down_block(sample) - - sample = self.conv_norm_out(sample) - sample = self.conv_act(sample) - sample = self.conv_out(sample) - - if self.latent_log_var == LogVarianceType.UNIFORM: - # Uniform Variance: model outputs N means and 1 shared log-variance channel. - # We need to expand the single logvar to match the number of means channels - # to create a format compatible with PER_CHANNEL (means + logvar, each with N channels). - # Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129) - # Target shape: (B, 2*N, ...) where first N are means, last N are logvar - - if sample.shape[1] < 2: - raise ValueError(f"Invalid channel count for UNIFORM mode: expected at least 2 channels (N means + 1 logvar), got {sample.shape[1]}") - - # Extract means (first N channels) and logvar (last 1 channel) - means = sample[:, :-1, ...] # (B, N, ...) - logvar = sample[:, -1:, ...] # (B, 1, ...) - - # Repeat logvar N times to match means channels - # Use expand/repeat pattern that works for both 4D and 5D tensors - num_channels = means.shape[1] - repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2) - repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...) - - # Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar] - sample = torch.cat([means, repeated_logvar], dim=1) - elif self.latent_log_var == LogVarianceType.CONSTANT: - sample = sample[:, :-1, ...] - approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects - sample = torch.cat( - [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0], - dim=1, - ) - - # Split into means and logvar, then normalize means - means, _ = torch.chunk(sample, 2, dim=1) - return self.per_channel_statistics.normalize(means) - - def tiled_encode( - self, - video: torch.Tensor, - tiling_config: TilingConfig | None = None, - ) -> torch.Tensor: - """Encode video to latent using tiled processing of the given video tensor. - Device Handling: - - Input video can be on CPU or GPU - - Accumulation buffers are created on model's device - - Each tile is automatically moved to model's device before encoding - - Output latent is returned on model's device - Args: - video: Input video tensor (B, 3, F, H, W) in range [-1, 1] - tiling_config: Tiling configuration for the video tensor - Returns: - Latent tensor (B, 128, F', H', W') on model's device - where F' = 1 + (F-1)/8, H' = H/32, W' = W/32 - """ - # Detect model device and dtype - model_device = next(self.parameters()).device - model_dtype = next(self.parameters()).dtype - - # Extract shape components - batch, _, frames, height, width = video.shape - - # Check frame count and crop if needed - if (frames - 1) % VIDEO_SCALE_FACTORS.time != 0: - frames_to_crop = (frames - 1) % VIDEO_SCALE_FACTORS.time - logger.warning(f"Number of frames {frames} of input video is not ({VIDEO_SCALE_FACTORS.time} * k + 1), last {frames_to_crop} frames will be cropped") - video = video[:, :, :-frames_to_crop, ...] - # Update frames after cropping - frames = video.shape[2] - - # Calculate output latent shape (inverse of upscale) - latent_shape = VideoLatentShape( - batch=batch, - channels=self.latent_channels, # 128 for standard VAE - frames=(frames - 1) // VIDEO_SCALE_FACTORS.time + 1, - height=height // VIDEO_SCALE_FACTORS.height, - width=width // VIDEO_SCALE_FACTORS.width, - ) - - # Prepare tiles (operates on VIDEO dimensions) - tiles = prepare_tiles_for_encoding(video, tiling_config) - - # Initialize accumulation buffers on model device - latent_buffer = torch.zeros( - latent_shape.to_torch_shape(), - device=model_device, - dtype=model_dtype, - ) - weights_buffer = torch.zeros_like(latent_buffer) - - # Process each tile - for tile in tiles: - # Extract video tile from input (may be on CPU) - video_tile = video[tile.in_coords] - - # Move tile to model device if needed - if video_tile.device != model_device or video_tile.dtype != model_dtype: - video_tile = video_tile.to(device=model_device, dtype=model_dtype) - - # Encode tile to latent (output on model device) - latent_tile = self.forward(video_tile) - - # Move blend mask to model device - mask = tile.blend_mask.to( - device=model_device, - dtype=model_dtype, - ) - - # Weighted accumulation in latent space - latent_buffer[tile.out_coords] += latent_tile * mask - weights_buffer[tile.out_coords] += mask - - del latent_tile, mask, video_tile - - # Normalize by accumulated weights - weights_buffer = weights_buffer.clamp(min=1e-8) - return latent_buffer / weights_buffer - - -def prepare_tiles_for_encoding( - video: torch.Tensor, - tiling_config: TilingConfig | None = None, -) -> List[Tile]: - """Prepare tiles for VAE encoding. - Args: - video: Input video tensor (B, 3, F, H, W) in range [-1, 1] - tiling_config: Tiling configuration for the video tensor - Returns: - List of tiles for the video tensor - """ - - splitters = [DEFAULT_SPLIT_OPERATION] * len(video.shape) - mappers = [DEFAULT_MAPPING_OPERATION] * len(video.shape) - minimum_spatial_overlap_px = 64 - minimum_temporal_overlap_frames = 16 - - if tiling_config is not None and tiling_config.spatial_config is not None: - cfg = tiling_config.spatial_config - - tile_size_px = cfg.tile_size_in_pixels - overlap_px = cfg.tile_overlap_in_pixels - - # Set minimum spatial overlap to 64 pixels in order to allow cutting padding from - # the front and back of the tiles and concatenate tiles without artifacts. - # The encoder uses symmetric padding (pad=1) in H and W at each conv layer. At tile - # boundaries, convs see padding (zeros/reflect) instead of real neighbor pixels, causing - # incorrect context near edges. - # For each overlap we discard 1 latent per edge (32px at scale 32) and concatenate tiles at a - # shared region with the next tile. - if overlap_px < minimum_spatial_overlap_px: - logger.warning( - f"Overlap pixels {overlap_px} in spatial tiling is less than \ - {minimum_spatial_overlap_px}, setting to minimum required {minimum_spatial_overlap_px}" - ) - overlap_px = minimum_spatial_overlap_px - - # Define split and map operations for the spatial dimensions - - # Height axis (H) - splitters[3] = split_in_spatial(tile_size_px, overlap_px) - mappers[3] = to_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height) - - # Width axis (W) - splitters[4] = split_in_spatial(tile_size_px, overlap_px) - mappers[4] = to_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width) - - if tiling_config is not None and tiling_config.temporal_config is not None: - cfg = tiling_config.temporal_config - tile_size_frames = cfg.tile_size_in_frames - overlap_frames = cfg.tile_overlap_in_frames - - if overlap_frames < minimum_temporal_overlap_frames: - logger.warning(f"Overlap frames {overlap_frames} is less than 16, setting to minimum required 16") - overlap_frames = minimum_temporal_overlap_frames - - splitters[2] = split_temporal(tile_size_frames, overlap_frames) - mappers[2] = to_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time) - - return create_tiles(video.shape, splitters, mappers) - - -def _make_decoder_block( - block_name: str, - block_config: dict[str, Any], - in_channels: int, - convolution_dimensions: int, - norm_layer: NormLayerType, - timestep_conditioning: bool, - norm_num_groups: int, - spatial_padding_mode: PaddingModeType, -) -> Tuple[nn.Module, int]: - out_channels = in_channels - if block_name == "res_x": - block = UNetMidBlock3D( - dims=convolution_dimensions, - in_channels=in_channels, - num_layers=block_config["num_layers"], - resnet_eps=1e-6, - resnet_groups=norm_num_groups, - norm_layer=norm_layer, - inject_noise=block_config.get("inject_noise", False), - timestep_conditioning=timestep_conditioning, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "attn_res_x": - block = UNetMidBlock3D( - dims=convolution_dimensions, - in_channels=in_channels, - num_layers=block_config["num_layers"], - resnet_groups=norm_num_groups, - norm_layer=norm_layer, - inject_noise=block_config.get("inject_noise", False), - timestep_conditioning=timestep_conditioning, - attention_head_dim=block_config["attention_head_dim"], - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "res_x_y": - out_channels = in_channels // block_config.get("multiplier", 2) - block = ResnetBlock3D( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=out_channels, - eps=1e-6, - groups=norm_num_groups, - norm_layer=norm_layer, - inject_noise=block_config.get("inject_noise", False), - timestep_conditioning=False, - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_time": - out_channels = in_channels // block_config.get("multiplier", 1) - block = DepthToSpaceUpsample( - dims=convolution_dimensions, - in_channels=in_channels, - stride=(2, 1, 1), - out_channels_reduction_factor=block_config.get("multiplier", 1), - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_space": - out_channels = in_channels // block_config.get("multiplier", 1) - block = DepthToSpaceUpsample( - dims=convolution_dimensions, - in_channels=in_channels, - stride=(1, 2, 2), - out_channels_reduction_factor=block_config.get("multiplier", 1), - spatial_padding_mode=spatial_padding_mode, - ) - elif block_name == "compress_all": - out_channels = in_channels // block_config.get("multiplier", 1) - block = DepthToSpaceUpsample( - dims=convolution_dimensions, - in_channels=in_channels, - stride=(2, 2, 2), - residual=block_config.get("residual", False), - out_channels_reduction_factor=block_config.get("multiplier", 1), - spatial_padding_mode=spatial_padding_mode, - ) - else: - raise ValueError(f"unknown layer: {block_name}") - - return block, out_channels - - -class VideoDecoder(nn.Module): - _DEFAULT_NORM_NUM_GROUPS = 32 - """ - Variational Autoencoder Decoder. Decodes latent representation into video frames. - The decoder upsamples latents through a series of upsampling operations (inverse of encoder). - Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration. - Upsampling blocks expand dimensions by 2x in specified dimensions: - - "compress_time": temporal only - - "compress_space": spatial only (H and W) - - "compress_all": all dimensions (F, H, W) - - "res_x" / "res_x_y" / "attn_res_x": no upsampling - Causal Mode: - causal=False (standard): Symmetric padding, allows future frame dependencies. - causal=True: Causal padding, each frame depends only on past/current frames. - First frame removed after temporal upsampling in both modes. Output shape unchanged. - Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes. - Args: - convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D). - in_channels: The number of input channels (latent channels). Default is 128. - out_channels: The number of output channels. For RGB images, this is 3. - decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params) - where params is either an int (num_layers) or a dict with configuration. - patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion: - H -> Hx4, W -> Wx4. Should be a power of 2. - norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`. - causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding. - When True, uses causal padding (past/current frames only). - timestep_conditioning: Whether to condition the decoder on timestep for denoising. - """ - - def __init__( - self, - convolution_dimensions: int = 3, - in_channels: int = 128, - out_channels: int = 3, - decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006 - patch_size: int = 4, - norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, - causal: bool = False, - timestep_conditioning: bool = False, - decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT, - base_channels: int = 128, - ): - super().__init__() - - # Spatiotemporal downscaling between decoded video space and VAE latents. - # According to the LTXV paper, the standard configuration downsamples - # video inputs by a factor of 8 in the temporal dimension and 32 in - # each spatial dimension (height and width). This parameter determines how - # many video frames and pixels correspond to a single latent cell. - self.video_downscale_factors = SpatioTemporalScaleFactors( - time=8, - height=32, - width=32, - ) - - self.patch_size = patch_size - out_channels = out_channels * patch_size**2 - self.causal = causal - self.timestep_conditioning = timestep_conditioning - self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS - - # Per-channel statistics for denormalizing latents - self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels) - - # Noise and timestep parameters for decoder conditioning - self.decode_noise_scale = 0.025 - self.decode_timestep = 0.05 - - # LTX VAE decoder architecture uses 3 upsampler blocks with multiplier equals to 2. - # Hence the total feature_channels is multiplied by 8 (2^3). - feature_channels = base_channels * 8 - - self.conv_in = make_conv_nd( - dims=convolution_dimensions, - in_channels=in_channels, - out_channels=feature_channels, - kernel_size=3, - stride=1, - padding=1, - causal=True, - spatial_padding_mode=decoder_spatial_padding_mode, - ) - - self.up_blocks = nn.ModuleList([]) - - for block_name, block_params in list(reversed(decoder_blocks)): - # Convert int to dict format for uniform handling - block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params - - block, feature_channels = _make_decoder_block( - block_name=block_name, - block_config=block_config, - in_channels=feature_channels, - convolution_dimensions=convolution_dimensions, - norm_layer=norm_layer, - timestep_conditioning=timestep_conditioning, - norm_num_groups=self._norm_num_groups, - spatial_padding_mode=decoder_spatial_padding_mode, - ) - - self.up_blocks.append(block) - - if norm_layer == NormLayerType.GROUP_NORM: - self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6) - elif norm_layer == NormLayerType.PIXEL_NORM: - self.conv_norm_out = PixelNorm() - - self.conv_act = nn.SiLU() - self.conv_out = make_conv_nd( - dims=convolution_dimensions, - in_channels=feature_channels, - out_channels=out_channels, - kernel_size=3, - padding=1, - causal=True, - spatial_padding_mode=decoder_spatial_padding_mode, - ) - - if timestep_conditioning: - self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0)) - self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim=feature_channels * 2, size_emb_dim=0) - self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels)) - - def forward( - self, - sample: torch.Tensor, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - r""" - Decode latent representation into video frames. - Args: - sample: Latent tensor (B, 128, F', H', W'). - timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None. - generator: Random generator for deterministic noise injection (if inject_noise=True in blocks). - Returns: - Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'. - Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512). - Note: First frame is removed after temporal upsampling regardless of causal mode. - When causal=False, allows future frame dependencies in convolutions but maintains same output shape. - """ - batch_size = sample.shape[0] - output_dtype = sample.dtype - weights_dtype = next(self.parameters()).dtype - sample = sample.to(weights_dtype) - - # Add noise if timestep conditioning is enabled - if self.timestep_conditioning: - noise = ( - torch.randn( - sample.size(), - generator=generator, - dtype=sample.dtype, - device=sample.device, - ) - * self.decode_noise_scale - ) - - sample = noise + (1.0 - self.decode_noise_scale) * sample - - # Denormalize latents - sample = self.per_channel_statistics.un_normalize(sample) - - # Use default decode_timestep if timestep not provided - if timestep is None and self.timestep_conditioning: - timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype) - - sample = self.conv_in(sample, causal=self.causal) - - scaled_timestep = None - if self.timestep_conditioning: - if timestep is None: - raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") - scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample) - - for up_block in self.up_blocks: - if isinstance(up_block, UNetMidBlock3D): - block_kwargs = { - "causal": self.causal, - "timestep": scaled_timestep if self.timestep_conditioning else None, - "generator": generator, - } - sample = up_block(sample, **block_kwargs) - elif isinstance(up_block, ResnetBlock3D): - sample = up_block(sample, causal=self.causal, generator=generator) - else: - sample = up_block(sample, causal=self.causal) - - sample = self.conv_norm_out(sample) - - if self.timestep_conditioning: - embedded_timestep = self.last_time_embedder( - timestep=scaled_timestep.flatten(), - hidden_dtype=sample.dtype, - ) - embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1) - ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(device=sample.device, dtype=sample.dtype) + embedded_timestep.reshape( - batch_size, - 2, - -1, - embedded_timestep.shape[-3], - embedded_timestep.shape[-2], - embedded_timestep.shape[-1], - ) - shift, scale = ada_values.unbind(dim=1) - sample = sample * (1 + scale) + shift - - sample = self.conv_act(sample) - sample = self.conv_out(sample, causal=self.causal) - - # Final spatial expansion: reverse the initial patchify from encoder - # Moves pixels from channels back to spatial dimensions - # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4 - sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1) - - return sample.to(output_dtype) - - def _prepare_tiles( - self, - latent: torch.Tensor, - tiling_config: TilingConfig | None = None, - ) -> List[Tile]: - splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape) - mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape) - if tiling_config is not None and tiling_config.spatial_config is not None: - cfg = tiling_config.spatial_config - long_side = max(latent.shape[3], latent.shape[4]) - - def enable_on_axis(axis_idx: int, factor: int) -> None: - size = cfg.tile_size_in_pixels // factor - overlap = cfg.tile_overlap_in_pixels // factor - axis_length = latent.shape[axis_idx] - lower_threshold = max(2, overlap + 1) - tile_size = max(lower_threshold, round(size * axis_length / long_side)) - splitters[axis_idx] = split_in_spatial(tile_size, overlap) - mappers[axis_idx] = to_mapping_operation(map_spatial_slice, scale=factor) - - enable_on_axis(3, self.video_downscale_factors.height) - enable_on_axis(4, self.video_downscale_factors.width) - - if tiling_config is not None and tiling_config.temporal_config is not None: - cfg = tiling_config.temporal_config - tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time - overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time - splitters[2] = split_in_temporal(tile_size, overlap) - mappers[2] = to_mapping_operation(map_temporal_slice, scale=self.video_downscale_factors.time) - - return create_tiles(latent.shape, splitters, mappers) - - def tiled_decode( - self, - latent: torch.Tensor, - tiling_config: TilingConfig | None = None, - timestep: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> Iterator[torch.Tensor]: - """ - Decode a latent tensor into video frames using tiled processing. - Splits the latent tensor into tiles, decodes each tile individually, - and yields video chunks as they become available. - Args: - latent: Input latent tensor (B, C, F', H', W'). - tiling_config: Tiling configuration for the latent tensor. - timestep: Optional timestep for decoder conditioning. - generator: Optional random generator for deterministic decoding. - Yields: - Video chunks (B, C, T, H, W) by temporal slices; - """ - - # Calculate full video shape from latent shape to get spatial dimensions - full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors) - tiles = self._prepare_tiles(latent, tiling_config) - - temporal_groups = self._group_tiles_by_temporal_slice(tiles) - - # State for temporal overlap handling - previous_chunk = None - previous_weights = None - previous_temporal_slice = None - - for temporal_group_tiles in temporal_groups: - curr_temporal_slice = temporal_group_tiles[0].out_coords[2] - - # Calculate the shape of the temporal buffer for this group of tiles. - # The temporal length depends on whether this is the first tile (starts at 0) or not. - # - First tile: (frames - 1) * scale + 1 - # - Subsequent tiles: frames * scale - # This logic is handled by TemporalAxisMapping and reflected in out_coords. - temporal_tile_buffer_shape = full_video_shape._replace( - frames=curr_temporal_slice.stop - curr_temporal_slice.start, - ) - - buffer = torch.zeros( - temporal_tile_buffer_shape.to_torch_shape(), - device=latent.device, - dtype=latent.dtype, - ) - - curr_weights = self._accumulate_temporal_group_into_buffer( - group_tiles=temporal_group_tiles, - buffer=buffer, - latent=latent, - timestep=timestep, - generator=generator, - ) - - # Blend with previous temporal chunk if it exists - if previous_chunk is not None: - # Check if current temporal slice overlaps with previous temporal slice - if previous_temporal_slice.stop > curr_temporal_slice.start: - overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start - temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None) - - # The overlap is already masked before it reaches this step. Each tile is accumulated into buffer - # with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add - # the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the - # previous buffers, then later normalize by weights. - previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :] - previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[:, :, slice(0, overlap_len), :, :] - - buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :] - curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[:, :, temporal_overlap_slice, :, :] - - # Yield the non-overlapping part of the previous chunk - previous_weights = previous_weights.clamp(min=1e-8) - yield_len = curr_temporal_slice.start - previous_temporal_slice.start - yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :] - - # Update state for next iteration - previous_chunk = buffer - previous_weights = curr_weights - previous_temporal_slice = curr_temporal_slice - - # Yield any remaining chunk - if previous_chunk is not None: - previous_weights = previous_weights.clamp(min=1e-8) - yield previous_chunk / previous_weights - - def decode_video( - self, - latent: torch.Tensor, - tiling_config: TilingConfig | None = None, - generator: torch.Generator | None = None, - ) -> Iterator[torch.Tensor]: - """Decode a video latent tensor, yielding float chunks ``[f, h, w, c]`` in ``[0, 1]``. - Subclasses (e.g. ``DistributedVideoDecoder``) may override this to - control eagerness or distribution across ranks. - """ - - def to_rgb(frames: torch.Tensor) -> torch.Tensor: - video = rearrange(frames[0], "c f h w -> f h w c") - return video.add_(1.0).mul_(0.5).clamp_(0.0, 1.0) - - if tiling_config is not None: - for frames in self.tiled_decode(latent, tiling_config, generator=generator): - yield to_rgb(frames) - else: - decoded = self(latent, generator=generator) - yield to_rgb(decoded) - - def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]: - """Group tiles by their temporal output slice.""" - if not tiles: - return [] - - groups = [] - current_slice = tiles[0].out_coords[2] - current_group = [] - - for tile in tiles: - tile_slice = tile.out_coords[2] - if tile_slice == current_slice: - current_group.append(tile) - else: - groups.append(current_group) - current_slice = tile_slice - current_group = [tile] - - # Add the final group - if current_group: - groups.append(current_group) - - return groups - - def _accumulate_temporal_group_into_buffer( - self, - group_tiles: List[Tile], - buffer: torch.Tensor, - latent: torch.Tensor, - timestep: torch.Tensor | None, - generator: torch.Generator | None, - ) -> torch.Tensor: - """ - Decode and accumulate all tiles of a temporal group into a local buffer. - The buffer is local to the group and always starts at time 0; temporal coordinates - are rebased by subtracting temporal_slice.start. - """ - temporal_slice = group_tiles[0].out_coords[2] - - weights = torch.zeros_like(buffer) - - for tile in group_tiles: - decoded_tile = self.forward(latent[tile.in_coords], timestep, generator) - mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype) - temporal_offset = tile.out_coords[2].start - temporal_slice.start - # Use the tile's output coordinate length, not the decoded tile's length, - # as the decoder may produce a different number of frames than expected - expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start - decoded_temporal_len = decoded_tile.shape[2] - - # Ensure we don't exceed the buffer or decoded tile bounds - actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset) - - chunk_coords = ( - slice(None), # batch - slice(None), # channels - slice(temporal_offset, temporal_offset + actual_temporal_len), - tile.out_coords[3], # height - tile.out_coords[4], # width - ) - - # Slice decoded_tile and mask to match the actual length we're writing - decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :] - mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask - - buffer[chunk_coords] += decoded_slice * mask_slice - weights[chunk_coords] += mask_slice - - return weights - - -def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int: - """ - Get the number of video chunks for a given number of frames and tiling configuration. - Args: - num_frames: Number of frames in the video. - tiling_config: Tiling configuration. - Returns: - Number of video chunks. - """ - if not tiling_config or not tiling_config.temporal_config: - return 1 - cfg = tiling_config.temporal_config - frame_stride = cfg.tile_size_in_frames - cfg.tile_overlap_in_frames - return (num_frames - 1 + frame_stride - 1) // frame_stride - - -def to_mapping_operation( - map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor | None]], - scale: int, -) -> MappingOperation: - """Create a mapping operation over a set of tiling intervals. - The given mapping function is applied to each interval in the input dimension. The result function is used for - creating tiles in the output dimension. - Args: - map_func: Mapping function to create the mapping operation from - scale: Scale factor for the transformation, used as an argument for the mapping function - Returns: - Mapping operation that takes a set of tiling intervals and returns a set of slices and masks in the output - dimension. - """ - - def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]: - output_slices: list[slice] = [] - masks_1d: list[torch.Tensor | None] = [] - for interval in intervals.intervals: - output_slice, mask_1d = map_func(interval.start, interval.end, interval.left_ramp, interval.right_ramp, scale) - output_slices.append(output_slice) - masks_1d.append(mask_1d) - return output_slices, masks_1d - - return map_op - - -def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]: - start = begin * scale - stop = 1 + (end - 1) * scale - left_ramp = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale - right_ramp = right_ramp * scale - - return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True) - - -def map_temporal_interval_to_latent(begin: int, end: int, left_ramp: int, right_ramp: int | None = None, scale: int = 1) -> Tuple[slice, torch.Tensor]: - """ - Map temporal interval in video frame space to latent space. - Args: - begin: Start position in video frame space - end: End position in video frame space - left_ramp: Left ramp size in video frame space - right_ramp: Right ramp size in video frame space - scale: Scale factor for transformation - Returns: - Tuple of (output_slice, blend_mask) - """ - start = begin // scale - stop = (end - 1) // scale + 1 - - left_ramp_latents = 0 if left_ramp == 0 else 1 + (left_ramp - 1) // scale - right_ramp_latents = right_ramp // scale - - if right_ramp_latents != 0: - raise ValueError("For tiled encoding, temporal tiles are expected to have a right ramp equal to 0") - - mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp_latents, right_ramp_latents) - - return slice(start, stop), mask_1d - - -def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]: - start = begin * scale - stop = end * scale - left_ramp = left_ramp * scale - right_ramp = right_ramp * scale - - return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False) - - -def map_spatial_interval_to_latent( - begin: int, - end: int, - left_ramp: int, - right_ramp: int, - scale: int, -) -> Tuple[slice, torch.Tensor]: - """Map spatial interval in pixel space to latent space. - Args: - begin: Start position in pixel space - end: End position in pixel space - left_ramp: Left ramp size in pixel space - right_ramp: Right ramp size in pixel space - scale: Scale factor for transformation - Returns: - Tuple of (output_slice, blend_mask) - """ - start = begin // scale - stop = end // scale - left_ramp = max(0, left_ramp // scale - 1) - - right_ramp = 0 if right_ramp == 0 else 1 - - mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp, right_ramp) - return slice(start, stop), mask_1d diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/__init__.py deleted file mode 100644 index 8bb6e0fe0..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -from ltx_core.quantization.fp8_cast import ( - TRANSFORMER_LINEAR_DOWNCAST_MAP, - UPCAST_DURING_INFERENCE, - UpcastWithStochasticRounding, - fp8_cast_fuse_rule, -) -from ltx_core.quantization.fp8_scaled_mm import fp8_scaled_mm_fuse_rule -from ltx_core.quantization.policy import QuantizationPolicy - -__all__ = [ - "TRANSFORMER_LINEAR_DOWNCAST_MAP", - "UPCAST_DURING_INFERENCE", - "QuantizationPolicy", - "UpcastWithStochasticRounding", - "fp8_cast_fuse_rule", - "fp8_scaled_mm_fuse_rule", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_cast.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_cast.py deleted file mode 100644 index c88b4b5fb..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_cast.py +++ /dev/null @@ -1,322 +0,0 @@ -from pathlib import Path - -import safetensors -import torch -from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule -from ltx_core.loader.kernels import TRITON_AVAILABLE -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import StateDict -from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps -from ltx_core.model.transformer.model import LTXModel -from ltx_core.quantization.policy import QuantizationPolicy - -BLOCK_SIZE = 1024 - - -def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor: - if not TRITON_AVAILABLE: - raise RuntimeError( - "fused_add_round_launch requires Triton, which is not available on this platform. " - "Callers should gate on ltx_core.loader.kernels.TRITON_AVAILABLE and use a " - "deterministic-rounding fallback instead." - ) - import triton # noqa: PLC0415 - from ltx_core.loader.kernels import fused_add_round_kernel # noqa: PLC0415 - - if original_weight.dtype == torch.float8_e4m3fn: - exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7 - elif original_weight.dtype == torch.float8_e5m2: - exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15 # noqa: F841 - else: - raise ValueError("Unsupported dtype") - - if target_weight.dtype != torch.bfloat16: - raise ValueError("target_weight dtype must be bfloat16") - - # Calculate grid and block sizes - n_elements = original_weight.numel() - grid = (triton.cdiv(n_elements, BLOCK_SIZE),) - - # Launch kernel - fused_add_round_kernel[grid]( - original_weight, - target_weight, - seed, - n_elements, - exponent_bias, - mantissa_bits, - BLOCK_SIZE, - ) - return target_weight - - -def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: - """ - Downcast the weight or bias to the float8_e4m3fn dtype. - """ - return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))] - - -def _upcast_and_round(weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0) -> torch.Tensor: - """ - Upcast the weight to the given dtype and optionally apply stochastic rounding. - Input weight needs to have float8_e4m3fn or float8_e5m2 dtype. - Stochastic rounding is implemented via a Triton kernel. When Triton is not - available (e.g., on Windows), this falls back to deterministic (nearest) - rounding via ``weight.to(dtype)``. - """ - if not with_stochastic_rounding or not TRITON_AVAILABLE or weight.device.type != "cuda": - return weight.to(dtype) - return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed) - - -class Fp8CastLinear(torch.nn.Linear): - """nn.Linear storing weights in fp8, upcasting to input dtype during forward. - Used via __class__ reassignment (not subclassing) so existing weight tensors - are preserved in-place. Class-level forward is required for torch.compile - compatibility — instance-level closure monkey-patches cause graph breaks. - """ - - _with_stochastic_rounding: bool - _seed: int - - def forward(self, input: torch.Tensor) -> torch.Tensor: # noqa: A002, type: ignore[override] - w_up = _upcast_and_round(self.weight, input.dtype, self._with_stochastic_rounding, self._seed) - b_up = _upcast_and_round(self.bias, input.dtype, self._with_stochastic_rounding, self._seed) if self.bias is not None else None - return torch.nn.functional.linear(input, w_up, b_up) - - -def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None: - """ - Intended to be applied via __class__ reassignment to existing nn.Linear - instances. Forward remains defined at the class level, which is required for - torch.compile compatibility — instance-level closure monkey-patches cause - graph breaks. - Also retypes ``weight`` and ``bias`` to fp8 so the meta param dtype matches - the post-load tensor dtype (sd_ops downcasts checkpoint bf16 -> fp8 at load). - Block streaming relies on this to derive pool buffer layout from the meta - model without an eager checkpoint read. - """ - layer.__class__ = Fp8CastLinear - layer._with_stochastic_rounding = with_stochastic_rounding - layer._seed = seed - layer.weight = torch.nn.Parameter( - torch.empty(layer.weight.shape, dtype=torch.float8_e4m3fn, device=layer.weight.device), - requires_grad=layer.weight.requires_grad, - ) - if layer.bias is not None: - layer.bias = torch.nn.Parameter( - torch.empty(layer.bias.shape, dtype=torch.float8_e4m3fn, device=layer.bias.device), - requires_grad=layer.bias.requires_grad, - ) - - -# Module-name suffixes for the Linears that participate in fp8 cast. Used by -# both the upcast matcher and the sd_ops downcast map so the two cannot drift. -# - ``.to_q`` / ``.to_k`` / ``.to_v`` / ``.to_out.0`` have a leading dot so they -# only match the attention Linears at ``...attnN.to_q`` etc. -# - ``ff.net.0.proj`` / ``ff.net.2`` are intentionally **dotless** so they match -# both video FF (``...ff.net.0.proj``) and audio FF (``...audio_ff.net.0.proj``). -_FP8_CAST_KEY_PREFIX = "transformer_blocks." -_FP8_CAST_LINEAR_SUFFIXES: tuple[str, ...] = ( - ".to_q", - ".to_k", - ".to_v", - ".to_out.0", - "ff.net.0.proj", - "ff.net.2", -) - - -def _is_fp8_cast_linear(module_name: str) -> bool: - """Return True if *module_name* names a Linear that should be fp8-cast.""" - if _FP8_CAST_KEY_PREFIX not in module_name: - return False - return any(module_name.endswith(suffix) for suffix in _FP8_CAST_LINEAR_SUFFIXES) - - -def _amend_forward_with_upcast(model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0) -> torch.nn.Module: - """ - Replace the forward method of the fp8-cast Linear layers (per - :data:`_FP8_CAST_LINEAR_SUFFIXES`) to forward with upcast and optional - stochastic rounding. - Only the Linears whose weights are downcast by :data:`TRANSFORMER_LINEAR_DOWNCAST_MAP` - are retyped. Linears outside that subset (e.g. ``to_gate_logits``) are left as - plain ``nn.Linear`` so the meta-model param dtype matches the loaded checkpoint - dtype. - """ - for name, m in model.named_modules(): - if isinstance(m, torch.nn.Linear) and _is_fp8_cast_linear(name): - _replace_fwd_with_upcast(m, with_stochastic_rounding, seed) - return model - - -def _build_transformer_linear_downcast_map() -> SDOps: - """Build the sd_ops downcast map from the same suffix registry as the matcher.""" - ops = SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP") - for suffix in _FP8_CAST_LINEAR_SUFFIXES: - ops = ops.with_kv_operation( - key_prefix=_FP8_CAST_KEY_PREFIX, - key_suffix=suffix + ".weight", - operation=_naive_weight_or_bias_downcast, - ).with_kv_operation( - key_prefix=_FP8_CAST_KEY_PREFIX, - key_suffix=suffix + ".bias", - operation=_naive_weight_or_bias_downcast, - ) - return ops - - -TRANSFORMER_LINEAR_DOWNCAST_MAP = _build_transformer_linear_downcast_map() - -UPCAST_DURING_INFERENCE = ModuleOps( - name="upcast_fp8_during_linear_forward", - matcher=lambda model: isinstance(model, LTXModel), - mutator=lambda model: _amend_forward_with_upcast(model, False), -) - - -class UpcastWithStochasticRounding(ModuleOps): - """ - ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype - and applying stochastic rounding during linear forward. - """ - - def __new__(cls, seed: int = 0): - return super().__new__( - cls, - name="upcast_fp8_during_linear_forward_with_stochastic_rounding", - matcher=lambda model: isinstance(model, LTXModel), - mutator=lambda model: _amend_forward_with_upcast(model, True, seed), - ) - - -def fuse_cast_fp8_weight( - delta_bf16: torch.Tensor, - weight_fp8: torch.Tensor, -) -> torch.Tensor: - """Return ``(delta_bf16 + dequantize(weight_fp8)).to(weight_fp8.dtype)``. - CUDA with Triton uses stochastic rounding via the fused kernel; otherwise - falls back to a deterministic bf16 add. ``delta_bf16`` is the bf16 - accumulator and is mutated in place. - """ - if delta_bf16.dtype != torch.bfloat16: - raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}") - if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE: - fused_add_round_launch(delta_bf16, weight_fp8, seed=0) - else: - delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16)) - return delta_bf16.to(dtype=weight_fp8.dtype) - - -def _fp8_cast_fuse( - key: str, - weight: torch.Tensor, - deltas: torch.Tensor, - model_sd: StateDict, -) -> dict[str, torch.Tensor]: - """Cast the dequantized FP8 weight + BF16 deltas back to ``weight.dtype`` - (FP8) via the fused-add-round kernel on CUDA. - Only a subset of linears are FP8-downcast (see ``TRANSFORMER_LINEAR_DOWNCAST_MAP``); - LoRAs may also target layers left in BF16 (e.g. audio ``add_q/k/v_proj``, cross-modal - projections). For those, fall back to a plain BF16 fuse. - """ - if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2): - return bf16_fuse_rule(key, weight, deltas, model_sd) - return {key: fuse_cast_fp8_weight(deltas, weight)} - - -fp8_cast_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_cast_fuse) - - -# Raw safetensors storage prefix shared by every diffusion-transformer -# parameter (and every prequant `*_scale` sibling). Verified against -# ltx-2.3-22b-{dev,distilled}-fp8.safetensors: 2924/2924 and 2992/2992 of -# the scale keys start with this exact prefix. -_RAW_DIFFUSION_MODEL_PREFIX = "model.diffusion_model." - - -def _read_scales(checkpoint_path: str | Path) -> dict[str, torch.Tensor]: - """Return ``{post_rename_param_key: scale_tensor}`` for every prequant - ``*_scale`` sibling in *checkpoint_path*. - Keys are returned in the post-rename form the loader will pass to the - sd-op (e.g. ``transformer_blocks.0.attn1.to_q.weight``) -- the raw - ``model.diffusion_model.`` prefix and the ``_scale`` suffix are both - stripped. Catches both ``.weight_scale`` and ``.bias_scale``; the - latter is absent in the current LTX-2.3 prequant checkpoints but - accepted for forward compatibility. - """ - out: dict[str, torch.Tensor] = {} - with safetensors.safe_open(str(checkpoint_path), framework="pt", device="cpu") as h: - raw_keys = h.keys() - for k in raw_keys: - if not k.endswith("_scale"): - continue - if not k.startswith(_RAW_DIFFUSION_MODEL_PREFIX): - raise ValueError(f"Scale key {k!r} does not start with the expected raw prefix {_RAW_DIFFUSION_MODEL_PREFIX!r}") - param_key = k.removeprefix(_RAW_DIFFUSION_MODEL_PREFIX).removesuffix("_scale") - out[param_key] = h.get_tensor(k) - return out - - -def _build_prequant_fold_sd_ops(scales: dict[str, torch.Tensor]) -> SDOps: - """Build sd-ops that fold prequant ``*_scale`` siblings into their parent - tensor at load time. - *scales* is keyed by the **post-rename** param key (e.g. - ``transformer_blocks.0.attn1.to_q.weight``); see :func:`_read_scales`. - Four ``with_kv_operation`` entries (symmetric for ``.weight`` and ``.bias``): - * ``.weight`` / ``.bias`` -> if a sibling scale exists in *scales*, fold; - then delegate to ``TRANSFORMER_LINEAR_DOWNCAST_MAP`` (downcast covered - Linears, pass everything else through). Without a scale, delegate - directly. - * ``.weight_scale`` / ``.bias_scale`` -> drop (the scale is consumed by - the fold). Raises if the scale key doesn't correspond to a known - entry in *scales* -- that means the file shipped a scale we didn't - pre-register, which would silently desync the fold. - """ - - def _on_param(param_key: str, value: torch.Tensor) -> list[KeyValueOperationResult]: - scale = scales.get(param_key) - if scale is None: - return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, value) - scale = scale.to(device=value.device) - if scale.ndim != 0: - raise ValueError(f"Unsupported scale shape {tuple(scale.shape)} for {param_key}") - bf16 = (value.to(torch.float32) * scale).to(torch.bfloat16) - # Delegate the final fp8-vs-bf16 decision to the downcast map: Linears - # outside the fp8 subset (e.g. to_gate_logits) stay bf16 to match the - # plain nn.Linear that the upcast matcher leaves untouched. - return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, bf16) - - def _drop_scale(scale_key: str, _value: torch.Tensor) -> list[KeyValueOperationResult]: - param_key = scale_key.removesuffix("_scale") - if param_key not in scales: - raise ValueError(f"Scale key {scale_key!r} has no matching entry in the prequant scales dict; _read_scales and the loader's rename map have drifted") - return [] - - # Register the drop ops first so the dict-membership sanity check is the - # earliest sd-op that can fire on a scale key -- we crash on a stray scale - # before any silently mismatched fold has a chance to land in the state - # dict. Registration order is irrelevant for correctness (no overlap - # between matchers) but communicates intent. - return ( - SDOps("FP8_CAST_PREQUANT_AWARE") - .with_kv_operation(key_suffix=".weight_scale", operation=_drop_scale) - .with_kv_operation(key_suffix=".bias_scale", operation=_drop_scale) - .with_kv_operation(key_suffix=".weight", operation=_on_param) - .with_kv_operation(key_suffix=".bias", operation=_on_param) - ) - - -def build_policy(checkpoint_path: str | Path) -> QuantizationPolicy: - """FP8 casting with upcasting during inference. - *checkpoint_path* is required (mirroring ``fp8_scaled_mm.build_policy``). - For prequantized fp8 checkpoints, sibling ``*_scale`` tensors (weight or - bias) are folded into the parent at load time. - """ - scales = _read_scales(checkpoint_path) - return QuantizationPolicy( - sd_ops=_build_prequant_fold_sd_ops(scales), - module_ops=(UPCAST_DURING_INFERENCE,), - fuse_rule=fp8_cast_fuse_rule, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_scaled_mm.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_scaled_mm.py deleted file mode 100644 index 23a3bef17..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_scaled_mm.py +++ /dev/null @@ -1,209 +0,0 @@ -import json -import struct -from typing import Callable - -import torch -from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.primitives import StateDict -from ltx_core.model.transformer import LTXModel -from ltx_core.quantization.policy import QuantizationPolicy -from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable -from torch import nn - - -def _read_safetensors_dtypes(path: str) -> dict[str, str]: - """Return ``{tensor_name: dtype_string}`` from the safetensors header.""" - with open(path, "rb") as f: - header_size = struct.unpack(" torch.Tensor: - origin_shape = x.shape - - if trtllm_scaled_mm_usable(): - qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale) - if qinput.dim() == 3: - qinput = qinput.reshape(-1, qinput.shape[-1]) - output = torch.ops.trtllm.cublas_scaled_mm( - qinput, - self.weight.t(), - scale_a=cur_input_scale, - scale_b=self.weight_scale, - bias=None, - out_dtype=x.dtype, - ) - else: - # Clamp before cast: out-of-range values cast to NaN/saturated FP8, which - # produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8). - fp8_min = torch.finfo(torch.float8_e4m3fn).min - fp8_max = torch.finfo(torch.float8_e4m3fn).max - qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn) - if qinput.dim() == 3: - qinput = qinput.reshape(-1, qinput.shape[-1]) - output = torch._scaled_mm( - qinput, - self.weight.t(), - scale_a=self.input_scale, - scale_b=self.weight_scale, - out_dtype=x.dtype, - use_fast_accum=True, - ) - - if self.bias is not None: - output = output + self.bias.to(output.dtype) - - if output.dim() != len(origin_shape): - output_shape = list(origin_shape) - output_shape[-1] = output.shape[-1] - output = output.reshape(output_shape) - - return output - - -def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize a weight tensor to ``float8_e4m3fn`` with a per-tensor scale.""" - weight_fp32 = weight.to(torch.float32) - - fp8_min = torch.finfo(torch.float8_e4m3fn).min - fp8_max = torch.finfo(torch.float8_e4m3fn).max - - max_abs = torch.amax(torch.abs(weight_fp32)) - scale = fp8_max / max_abs - - @torch.compiler.disable - def _quantize(weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - weight_scale = scale.reciprocal() - return quantized_weight, weight_scale - - quantized_weight, weight_scale = _quantize(weight_fp32, scale, fp8_min, fp8_max) - return quantized_weight, weight_scale - - -def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear: - """Create an ``FP8Linear`` matching the shape/bias of *layer*.""" - return FP8Linear( - in_features=layer.in_features, - out_features=layer.out_features, - bias=layer.bias is not None, - device=layer.weight.device, - ) - - -def _swap_linears_to_fp8(model: nn.Module, should_swap: Callable[[str], bool]) -> nn.Module: - """Replace nn.Linear layers with FP8Linear where ``should_swap(name)`` returns True.""" - replacements: list[tuple[nn.Module, str, nn.Linear]] = [] - - for name, module in model.named_modules(): - if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear): - continue - if not should_swap(name): - continue - - if "." in name: - parent_name, attr_name = name.rsplit(".", 1) - parent = model.get_submodule(parent_name) - else: - parent = model - attr_name = name - - replacements.append((parent, attr_name, module)) - - for parent, attr_name, linear in replacements: - setattr(parent, attr_name, _linear_to_fp8linear(linear)) - - return model - - -def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]: - """Return the FP8 swap ``ModuleOps`` for layers whose ``.weight`` is ``F8_E4M3`` - and which have a sibling ``.weight_scale`` tensor in the checkpoint. - Raises ``ValueError`` if no such layers are found — that combination is ambiguous - (a BF16 checkpoint with this policy would load as a no-op). - """ - dtypes = _read_safetensors_dtypes(checkpoint_path) - fp8_scale_paths = frozenset(key.removesuffix(".weight_scale") for key in dtypes if key.endswith(".weight_scale") and dtypes.get(key.removesuffix(".weight_scale") + ".weight") == "F8_E4M3") - if not fp8_scale_paths: - raise ValueError( - f"fp8_scaled_mm requires a pre-quantized checkpoint with F8_E4M3 .weight + .weight_scale tensors, but {checkpoint_path!r} has none. Use QuantizationPolicy.fp8_cast() for BF16 checkpoints." - ) - - def _should_swap(name: str) -> bool: - suffix = "." + name - return any(p == name or p.endswith(suffix) for p in fp8_scale_paths) - - return ( - ModuleOps( - name="fp8_swap_linears", - matcher=lambda model: isinstance(model, LTXModel), - mutator=lambda model: _swap_linears_to_fp8(model, _should_swap), - ), - ) - - -def _fp8_scaled_mm_fuse( - key: str, - weight: torch.Tensor, - deltas: torch.Tensor, - model_sd: StateDict, -) -> dict[str, torch.Tensor]: - """Dequantize via ``weight.float() * weight_scale``, add the BF16 delta, - and re-quantize to FP8 with a fresh per-tensor scale. - Layers that were not swapped to scaled FP8 (e.g. small embedder linears - excluded from the auto-discovered swap set) stay BF16 and have no - ``.weight_scale`` companion -- for those, fall back to a plain bf16 fuse. - """ - scale_key = key.replace(".weight", ".weight_scale") - if scale_key not in model_sd.sd: - return bf16_fuse_rule(key, weight, deltas, model_sd) - weight_scale = model_sd.sd[scale_key] - original_weight = weight.to(torch.float32) * weight_scale - new_weight = original_weight + deltas.to(torch.float32) - new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight) - return {key: new_fp8_weight, scale_key: new_weight_scale} - - -fp8_scaled_mm_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_scaled_mm_fuse) - - -def build_policy(checkpoint_path: str) -> QuantizationPolicy: - """FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales. - The set of layers to swap to ``FP8Linear`` is discovered from the - checkpoint's ``.weight_scale`` tensors via suffix-matching against the - model's named modules. Requires a pre-quantized checkpoint; for BF16 - checkpoints, use :func:`ltx_core.quantization.fp8_cast.build_policy`. - """ - return QuantizationPolicy( - sd_ops=None, - module_ops=get_fp8_swap_module_ops(checkpoint_path), - fuse_rule=fp8_scaled_mm_fuse_rule, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/policy.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/policy.py deleted file mode 100644 index 20e991b25..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/policy.py +++ /dev/null @@ -1,24 +0,0 @@ -from dataclasses import dataclass - -from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.model.transformer.model import LTXModel - - -@dataclass(frozen=True) -class QuantizationPolicy: - """Configuration for model quantization during loading. - Attributes: - sd_ops: State-dict operations applied to each tensor during load. - module_ops: Post-load module transformations applied to the meta model. - model_configurator: Configurator class to use when constructing the transformer. - fuse_rule: How LoRA deltas merge into this policy's weight layout. - Default ``bf16_fuse_rule`` is used when no policy is configured. - """ - - sd_ops: SDOps | None = None - module_ops: tuple[ModuleOps, ...] = () - model_configurator: type[ModelConfigurator[LTXModel]] | None = None - fuse_rule: FuseRule = bf16_fuse_rule diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/trtllm_scaled_usable.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/trtllm_scaled_usable.py deleted file mode 100644 index deb69d03e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/trtllm_scaled_usable.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Runtime detection of TensorRT-LLM FP8 scaled-matmul availability. -When the TRT-LLM ops are usable on the current host (Linux + Hopper-class CUDA -+ tensorrt_llm wheel installed) we use them since they outperform the PyTorch-native -``torch._scaled_mm`` path. Otherwise we fall back to the native implementation, -which is portable across platforms (Windows, macOS, AMD GPUs). -The check runs once and is cached. -""" - -from __future__ import annotations - -import platform -from functools import cache - -import torch - - -@cache -def trtllm_scaled_mm_usable() -> bool: - if platform.system() != "Linux": - return False - - if not torch.cuda.is_available(): - return False - - major, minor = torch.cuda.get_device_capability() - sm = major * 10 + minor - - if sm < 90 or sm >= 120: - return False - - # The import is load-bearing — registers the trtllm torch ops as a side effect. - try: - import tensorrt_llm # noqa: F401, PLC0415 - except Exception: - return False - - return True diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/__init__.py deleted file mode 100644 index b65ae143a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""CLIP/text encoder model components.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/__init__.py deleted file mode 100644 index 009a38d4b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Gemma text encoder components.""" - -from ltx_core.text_encoders.gemma.embeddings_processor import ( - EmbeddingsProcessor, - EmbeddingsProcessorOutput, - convert_to_additive_mask, -) -from ltx_core.text_encoders.gemma.encoders.base_encoder import ( - GemmaTextEncoder, - module_ops_from_gemma_root, -) -from ltx_core.text_encoders.gemma.encoders.encoder_configurator import ( - EMBEDDINGS_PROCESSOR_KEY_OPS, - GEMMA_LLM_KEY_OPS, - GEMMA_MODEL_OPS, - VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS, - EmbeddingsProcessorConfigurator, - GemmaTextEncoderConfigurator, -) - -__all__ = [ - "EMBEDDINGS_PROCESSOR_KEY_OPS", - "GEMMA_LLM_KEY_OPS", - "GEMMA_MODEL_OPS", - "VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS", - "EmbeddingsProcessor", - "EmbeddingsProcessorConfigurator", - "EmbeddingsProcessorOutput", - "GemmaTextEncoder", - "GemmaTextEncoderConfigurator", - "convert_to_additive_mask", - "module_ops_from_gemma_root", -] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/config.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/config.py deleted file mode 100644 index 8b23e91b0..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/config.py +++ /dev/null @@ -1,75 +0,0 @@ -from dataclasses import asdict, dataclass, field - - -@dataclass -class Gemma3RopeScaling: - factor: float = 8.0 - rope_type: str = "linear" - - -@dataclass -class Gemma3TextConfig: - attention_bias: bool = False - attention_dropout: float = 0.0 - attn_logit_softcapping: float | None = None - cache_implementation: str = "hybrid" - final_logit_softcapping: float | None = None - head_dim: int = 256 - hidden_activation: str = "gelu_pytorch_tanh" - hidden_size: int = 3840 - initializer_range: float = 0.02 - intermediate_size: int = 15360 - max_position_embeddings: int = 131072 - model_type: str = "gemma3_text" - num_attention_heads: int = 16 - num_hidden_layers: int = 48 - num_key_value_heads: int = 8 - query_pre_attn_scalar: int = 256 - rms_norm_eps: float = 1e-06 - rope_local_base_freq: int = 10000 - rope_scaling: Gemma3RopeScaling = field(default_factory=Gemma3RopeScaling) - rope_theta: int = 1000000 - sliding_window: int = 1024 - sliding_window_pattern: int = 6 - torch_dtype: str = "float32" - use_cache: bool = True - vocab_size: int = 262208 - - -@dataclass -class Gemma3VisionConfig: - attention_dropout: float = 0.0 - hidden_act: str = "gelu_pytorch_tanh" - hidden_size: int = 1152 - image_size: int = 896 - intermediate_size: int = 4304 - layer_norm_eps: float = 1e-06 - model_type: str = "siglip_vision_model" - num_attention_heads: int = 16 - num_channels: int = 3 - num_hidden_layers: int = 27 - patch_size: int = 14 - torch_dtype: str = "float32" - vision_use_head: bool = False - - -@dataclass -class Gemma3ConfigData: - architectures: list[str] = field(default_factory=lambda: ["Gemma3ForConditionalGeneration"]) - boi_token_index: int = 255999 - eoi_token_index: int = 256000 - eos_token_id: list[int] = field(default_factory=lambda: [1, 106]) - image_token_index: int = 262144 - initializer_range: float = 0.02 - mm_tokens_per_image: int = 256 - model_type: str = "gemma3" - text_config: Gemma3TextConfig = field(default_factory=Gemma3TextConfig) - torch_dtype: str = "bfloat16" - transformers_version: str = "4.51.0" - vision_config: Gemma3VisionConfig = field(default_factory=Gemma3VisionConfig) - - def to_dict(self) -> dict: - return asdict(self) - - -GEMMA3_CONFIG_FOR_LTX = Gemma3ConfigData() diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_connector.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_connector.py deleted file mode 100644 index 3239224a4..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_connector.py +++ /dev/null @@ -1,241 +0,0 @@ -import torch -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.model.transformer.attention import Attention -from ltx_core.model.transformer.feed_forward import FeedForward -from ltx_core.model.transformer.rope import ( - LTXRopeType, - generate_freq_grid_np, - generate_freq_grid_pytorch, - precompute_freqs_cis, -) -from ltx_core.utils import rms_norm - - -class _BasicTransformerBlock1D(torch.nn.Module): - def __init__( - self, - dim: int, - heads: int, - dim_head: int, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - apply_gated_attention: bool = False, - ): - super().__init__() - - self.attn1 = Attention( - query_dim=dim, - heads=heads, - dim_head=dim_head, - rope_type=rope_type, - apply_gated_attention=apply_gated_attention, - ) - - self.ff = FeedForward( - dim, - dim_out=dim, - ) - - def forward( - self, - hidden_states: torch.Tensor, - additive_attention_mask: torch.Tensor | None = None, - pe: torch.Tensor | None = None, - ) -> torch.Tensor: - # Notice that normalization is always applied before the real computation in the following blocks. - - # 1. Normalization Before Self-Attention - norm_hidden_states = rms_norm(hidden_states) - - norm_hidden_states = norm_hidden_states.squeeze(1) - - # 2. Self-Attention — `mask` is the kernel-boundary name for the additive mask. - attn_output = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe) - - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # 3. Normalization before Feed-Forward - norm_hidden_states = rms_norm(hidden_states) - - # 4. Feed-forward - ff_output = self.ff(norm_hidden_states) - - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states - - -class Embeddings1DConnector(torch.nn.Module): - """ - Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or - other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can - substitute padded positions with learnable registers. The module is highly configurable for head size, number of - layers, and register usage. - Args: - attention_head_dim (int): Dimension of each attention head (default=128). - num_attention_heads (int): Number of attention heads (default=30). - num_layers (int): Number of transformer layers (default=2). - positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0). - positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]). - causal_temporal_positioning (bool): If True, uses causal attention (default=False). - num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables - register replacement. (default=128) - rope_type (LTXRopeType): The RoPE variant to use. - double_precision_rope (bool): Use double precision rope calculation (default=False). - """ - - _supports_gradient_checkpointing = True - - def __init__( - self, - attention_head_dim: int = 128, - num_attention_heads: int = 30, - num_layers: int = 2, - positional_embedding_theta: float = 10000.0, - positional_embedding_max_pos: list[int] | None = None, - causal_temporal_positioning: bool = False, - num_learnable_registers: int | None = 128, - rope_type: LTXRopeType = LTXRopeType.SPLIT, - double_precision_rope: bool = False, - apply_gated_attention: bool = False, - ): - super().__init__() - self.num_attention_heads = num_attention_heads - self.inner_dim = num_attention_heads * attention_head_dim - self.causal_temporal_positioning = causal_temporal_positioning - self.positional_embedding_theta = positional_embedding_theta - self.positional_embedding_max_pos = positional_embedding_max_pos if positional_embedding_max_pos is not None else [1] - self.rope_type = rope_type - self.double_precision_rope = double_precision_rope - self.transformer_1d_blocks = torch.nn.ModuleList( - [ - _BasicTransformerBlock1D( - dim=self.inner_dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - rope_type=rope_type, - apply_gated_attention=apply_gated_attention, - ) - for _ in range(num_layers) - ] - ) - - self.num_learnable_registers = num_learnable_registers - if self.num_learnable_registers: - self.learnable_registers = torch.nn.Parameter(torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0) - - def _replace_padded_with_learnable_registers(self, hidden_states: torch.Tensor, additive_attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - batch_size, seq_len, _ = hidden_states.shape - - assert seq_len % self.num_learnable_registers == 0 - - registers = self.learnable_registers.repeat(seq_len // self.num_learnable_registers, 1).to(hidden_states.dtype) - registers = registers.unsqueeze(0).expand(batch_size, -1, -1) # (B, seq_len, hidden_dim) - binary_mask = additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0 - binary_mask = binary_mask.to(hidden_states.dtype) - hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers - - return hidden_states, torch.zeros_like(additive_attention_mask) - - def forward( - self, - hidden_states: torch.Tensor, - additive_attention_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Forward pass of Embeddings1DConnector. - Args: - hidden_states: (B, S, D) input embeddings. - additive_attention_mask: optional additive mask of shape (B, 1, 1, S), where - valid = 0.0 and padding = -torch.finfo(dtype).max. - Returns: - (hidden_states, additive_attention_mask) - """ - if self.num_learnable_registers: - hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers(hidden_states, additive_attention_mask) - - indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device) - indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1) - freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch - freqs_cis = precompute_freqs_cis( - indices_grid=indices_grid, - dim=self.inner_dim, - out_dtype=hidden_states.dtype, - theta=self.positional_embedding_theta, - max_pos=self.positional_embedding_max_pos, - num_attention_heads=self.num_attention_heads, - rope_type=self.rope_type, - freq_grid_generator=freq_grid_generator, - ) - - for block in self.transformer_1d_blocks: - hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis) - - hidden_states = rms_norm(hidden_states) - - return hidden_states, additive_attention_mask - - -class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]): - """Configurator for video embeddings connector.""" - - @classmethod - def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector: - transformer_config = config.get("transformer", {}) - rope_type = LTXRopeType(transformer_config.get("rope_type", "split")) - double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64" - pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1]) - - # Video connector dimensions - num_attention_heads = transformer_config.get("connector_num_attention_heads", 30) - attention_head_dim = transformer_config.get("connector_attention_head_dim", 128) - num_layers = transformer_config.get("connector_num_layers", 2) - - connector = Embeddings1DConnector( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_layers=num_layers, - positional_embedding_max_pos=pe_max_pos, - rope_type=rope_type, - double_precision_rope=double_precision_rope, - apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False), - ) - return connector - - -class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]): - """Configurator for audio embeddings connector with separate dimension settings.""" - - @classmethod - def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector: - transformer_config = config.get("transformer", {}) - rope_type = LTXRopeType(transformer_config.get("rope_type", "split")) - double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64" - pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1]) - - # Audio connector dimensions - fall back to video connector config for backwards compatibility - num_attention_heads = transformer_config.get( - "audio_connector_num_attention_heads", - transformer_config.get("connector_num_attention_heads", 30), - ) - attention_head_dim = transformer_config.get( - "audio_connector_attention_head_dim", - transformer_config.get("connector_attention_head_dim", 128), - ) - num_layers = transformer_config.get( - "audio_connector_num_layers", - transformer_config.get("connector_num_layers", 2), - ) - - connector = Embeddings1DConnector( - num_attention_heads=num_attention_heads, - attention_head_dim=attention_head_dim, - num_layers=num_layers, - positional_embedding_max_pos=pe_max_pos, - rope_type=rope_type, - double_precision_rope=double_precision_rope, - apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False), - ) - return connector diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_processor.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_processor.py deleted file mode 100644 index e54804e7b..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_processor.py +++ /dev/null @@ -1,113 +0,0 @@ -from typing import NamedTuple - -import torch -from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnector -from torch import nn - - -class EmbeddingsProcessorOutput(NamedTuple): - video_encoding: torch.Tensor - audio_encoding: torch.Tensor | None - attention_mask: torch.Tensor - - -def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Convert binary attention mask to additive form for transformer masking.""" - return (attention_mask.to(torch.int64) - 1).to(dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(dtype).max - - -def _compute_right_pad_order(additive_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Compute the index permutation that places valid tokens before pads in each row. - Stable sort: valid tokens keep their relative order. Idempotent for inputs already - right-padded. The sort and reordered mask depend only on the mask, so they can be - computed once and reused across multiple feature tensors that share the mask. - Args: - additive_mask: (B, 1, 1, S) additive mask, ``0.0`` for valid, ``-finfo.max`` for pad. - Returns: - ``(sort_idx, reordered_additive_mask)``: ``sort_idx`` is (B, S); the reordered mask - has the same shape as the input. - """ - binary = (additive_mask[:, 0, 0, :] >= 0).to(torch.int32) # (B, S) - sort_idx = torch.argsort(binary, dim=-1, descending=True, stable=True) # (B, S) - new_binary = torch.gather(binary, 1, sort_idx) - new_additive = (new_binary.to(additive_mask.dtype) - 1) * torch.finfo(additive_mask.dtype).max - return sort_idx, new_additive[:, None, None, :] - - -def _apply_right_pad_order(features: torch.Tensor, sort_idx: torch.Tensor) -> torch.Tensor: - """Apply a precomputed right-pad permutation (from ``_compute_right_pad_order``) to features.""" - return torch.gather(features, 1, sort_idx.unsqueeze(-1).expand_as(features)) - - -def _to_binary_mask(encoded_mask: torch.Tensor, lead_shape: tuple[int, int]) -> torch.Tensor: - """Convert connector output mask to a binary (0/1) mask shaped ``(B, S, 1)`` for broadcasting.""" - return (encoded_mask < 0.000001).to(torch.int64).reshape([lead_shape[0], lead_shape[1], 1]) - - -class EmbeddingsProcessor(nn.Module): - """Wraps feature extractor + video connector + optional audio connector. - Can operate in two modes: - 1. create_embeddings(): Takes pre-computed features + additive mask (backward compat, used by trainer) - 2. process_hidden_states(): Takes raw Gemma hidden states, runs feature extraction + connectors - """ - - def __init__( - self, - *, - feature_extractor: nn.Module | None = None, - video_connector: Embeddings1DConnector, - audio_connector: Embeddings1DConnector | None = None, - ): - super().__init__() - self.feature_extractor = feature_extractor - self.video_connector = video_connector - self.audio_connector = audio_connector - - def create_embeddings( - self, - video_features: torch.Tensor, - audio_features: torch.Tensor | None, - additive_attention_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: - if self.audio_connector is not None and audio_features is None: - raise ValueError("Audio connector is configured but no audio features were provided.") - if self.audio_connector is None and audio_features is not None: - raise ValueError("Audio features were provided but no audio connector is configured.") - - # Connectors expect right-padded input ([valid, pad]). Normalize layout here so the - # upstream tokenizer can keep using either side without coupling to the connector. - # The sort index depends only on the mask, so compute it once and reuse for audio. - sort_idx, mask_for_connector = _compute_right_pad_order(additive_attention_mask) - video_features = _apply_right_pad_order(video_features, sort_idx) - video_encoded, video_mask = self.video_connector(video_features, mask_for_connector) - binary_mask = _to_binary_mask(video_mask, video_encoded.shape[:2]) - video_encoded = video_encoded * binary_mask - - audio_encoded = None - if self.audio_connector is not None: - audio_features = _apply_right_pad_order(audio_features, sort_idx) - audio_encoded, _ = self.audio_connector(audio_features, mask_for_connector) - - return video_encoded, audio_encoded, binary_mask.squeeze(-1) - - def process_hidden_states( - self, - hidden_states: tuple[torch.Tensor, ...], - attention_mask: torch.Tensor, - padding_side: str = "left", - ) -> EmbeddingsProcessorOutput: - """Full pipeline: feature extraction -> connectors -> final embeddings. - Args: - hidden_states: Raw Gemma hidden states (tuple of tensors per layer). - attention_mask: Binary attention mask [B, seq_len]. - padding_side: Padding side used during tokenization. - Returns: - EmbeddingsProcessorOutput with video_encoding, audio_encoding, and attention_mask. - """ - if self.feature_extractor is None: - raise ValueError("feature_extractor is required for process_hidden_states()") - - video_feats, audio_feats = self.feature_extractor(hidden_states, attention_mask, padding_side) - additive_mask = convert_to_additive_mask(attention_mask, video_feats.dtype) - video_enc, audio_enc, binary_mask = self.create_embeddings(video_feats, audio_feats, additive_mask) - return EmbeddingsProcessorOutput(video_enc, audio_enc, binary_mask) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/base_encoder.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/base_encoder.py deleted file mode 100644 index 81168574c..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/base_encoder.py +++ /dev/null @@ -1,211 +0,0 @@ -import functools -from pathlib import Path - -import torch -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer -from ltx_core.utils import find_matching_file -from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor - - -class GemmaTextEncoder(torch.nn.Module): - """Pure Gemma text encoder — runs the LLM and returns raw hidden states. - Prompt enhancement (generate) is also supported since the full - Gemma3ForConditionalGeneration model (including lm_head) is loaded. - """ - - def __init__( - self, - model: Gemma3ForConditionalGeneration | None = None, - tokenizer: LTXVGemmaTokenizer | None = None, - processor: Gemma3Processor | None = None, - dtype: torch.dtype = torch.bfloat16, - ): - super().__init__() - self.model = model - self.tokenizer = tokenizer - self.processor = processor - self._dtype = dtype - - def encode( - self, - prompts: list[str], - padding_side: str = "left", # noqa: ARG002 - ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]: - """Run a single fused Gemma forward over a batch of prompts. - Calls the inner model (self.model.model) to skip lm_head logits computation - (~500 MiB saving). The tokenizer pads every prompt to ``max_length`` (1024), - so the inputs stack into a single ``[N, 1024]`` batch with no further padding - logic; per-prompt outputs are sliced back to ``[1, 1024, D]`` / ``[1, 1024]`` - in the original order. - """ - if not prompts: - return [] - tokenized = [self.tokenizer.tokenize_with_weights(t)["gemma"] for t in prompts] - input_ids = torch.tensor( - [[tok for tok, _ in pairs] for pairs in tokenized], - device=self.model.device, - ) - attention_mask = torch.tensor( - [[w for _, w in pairs] for pairs in tokenized], - device=self.model.device, - ) - outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) - hidden_states = outputs.hidden_states - del outputs - return [(tuple(h[i : i + 1] for h in hidden_states), attention_mask[i : i + 1]) for i in range(len(prompts))] - - # --- Prompt enhancement methods --- - - def _enhance( - self, - messages: list[dict[str, str]], - image: torch.Tensor | None = None, - max_new_tokens: int = 512, - seed: int = 10, - ) -> str: - text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - - model_inputs = self.processor( - text=text, - images=image, - return_tensors="pt", - ).to(self.model.device) - pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0 - model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id) - - with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]): - torch.manual_seed(seed) - outputs = self.model.generate( - **model_inputs, - max_new_tokens=max_new_tokens, - do_sample=True, - temperature=0.7, - ) - generated_ids = outputs[0][len(model_inputs.input_ids[0]) :] - enhanced_prompt = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=True) - - return enhanced_prompt - - def enhance_t2v( - self, - prompt: str, - max_new_tokens: int = 512, - system_prompt: str | None = None, - seed: int = 10, - ) -> str: - """Enhance a text prompt for T2V generation.""" - system_prompt = system_prompt or self.default_gemma_t2v_system_prompt - - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": f"user prompt: {prompt}"}, - ] - - return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed) - - def enhance_i2v( - self, - prompt: str, - image: torch.Tensor, - max_new_tokens: int = 512, - system_prompt: str | None = None, - seed: int = 10, - ) -> str: - """Enhance a text prompt for I2V generation using a reference image.""" - system_prompt = system_prompt or self.default_gemma_i2v_system_prompt - messages = [ - {"role": "system", "content": system_prompt}, - { - "role": "user", - "content": [ - {"type": "image"}, - {"type": "text", "text": f"User Raw Input Prompt: {prompt}."}, - ], - }, - ] - return self._enhance(messages, image=image, max_new_tokens=max_new_tokens, seed=seed) - - @functools.cached_property - def default_gemma_i2v_system_prompt(self) -> str: - return _load_system_prompt("gemma_i2v_system_prompt.txt") - - @functools.cached_property - def default_gemma_t2v_system_prompt(self) -> str: - return _load_system_prompt("gemma_t2v_system_prompt.txt") - - -# --- Standalone utility functions --- - - -@functools.lru_cache(maxsize=2) -def _load_system_prompt(prompt_name: str) -> str: - with open(Path(__file__).parent / "prompts" / f"{prompt_name}", "r") as f: - return f.read() - - -def _cat_with_padding( - tensor: torch.Tensor, - padding_length: int, - value: int | float, -) -> torch.Tensor: - """Concatenate a tensor with a padding tensor of the given value.""" - return torch.cat( - [ - tensor, - torch.full( - (1, padding_length), - value, - dtype=tensor.dtype, - device=tensor.device, - ), - ], - dim=1, - ) - - -def _pad_inputs_for_attention_alignment( - model_inputs: dict[str, torch.Tensor], - pad_token_id: int = 0, - alignment: int = 8, -) -> dict[str, torch.Tensor]: - """Pad sequence length to multiple of alignment for Flash Attention compatibility.""" - seq_len = model_inputs.input_ids.shape[1] - padded_len = ((seq_len + alignment - 1) // alignment) * alignment - padding_length = padded_len - seq_len - - if padding_length > 0: - model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id) - model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0) - if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None: - model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0) - - return model_inputs - - -def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]: - tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent) - processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent) - - def load_tokenizer(module: GemmaTextEncoder) -> GemmaTextEncoder: - module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024) - return module - - def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder: - image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True, use_fast=False) - if not module.tokenizer: - raise ValueError("Tokenizer model operation must be performed before processor model operation") - module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer) - return module - - tokenizer_load_ops = ModuleOps( - "TokenizerLoad", - matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.tokenizer is None, - mutator=load_tokenizer, - ) - processor_load_ops = ModuleOps( - "ProcessorLoad", - matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.processor is None, - mutator=load_processor, - ) - return (tokenizer_load_ops, processor_load_ops) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py deleted file mode 100644 index 902b751a2..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py +++ /dev/null @@ -1,254 +0,0 @@ -import torch -import transformers -from ltx_core.loader import KeyValueOperationResult -from ltx_core.loader.module_ops import ModuleOps -from ltx_core.loader.sd_ops import SDOps -from ltx_core.model.model_protocol import ModelConfigurator -from ltx_core.text_encoders.gemma.config import GEMMA3_CONFIG_FOR_LTX -from ltx_core.text_encoders.gemma.embeddings_connector import ( - AudioEmbeddings1DConnectorConfigurator, - Embeddings1DConnectorConfigurator, -) -from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor -from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder -from ltx_core.text_encoders.gemma.feature_extractor import ( - FeatureExtractorV1, - FeatureExtractorV2, -) -from transformers import Gemma3Config -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS -from transformers.models.gemma3 import Gemma3ForConditionalGeneration - -_TRANSFORMERS_V5: bool = int(transformers.__version__.split(".", 1)[0]) >= 5 - - -class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]): - @classmethod - def from_config(cls, config: dict) -> GemmaTextEncoder: # noqa: ARG003 - gemma_config = Gemma3Config.from_dict(GEMMA3_CONFIG_FOR_LTX.to_dict()) - with torch.device("meta"): - model = Gemma3ForConditionalGeneration(gemma_config) - - return GemmaTextEncoder(model=model) - - -class EmbeddingsProcessorConfigurator(ModelConfigurator[EmbeddingsProcessor]): - @classmethod - def from_config(cls, config: dict) -> EmbeddingsProcessor: - transformer_config = config.get("transformer", {}) - - # Create video embeddings connector (always needed) - video_connector = Embeddings1DConnectorConfigurator.from_config(config) - - # Create audio embeddings connector - audio_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config) - - # Create feature extractor - feature_extractor = _create_feature_extractor(transformer_config) - - return EmbeddingsProcessor( - video_connector=video_connector, - audio_connector=audio_connector, - feature_extractor=feature_extractor, - ) - - -_V2_EXPECTED_CONFIG = { - "caption_proj_before_connector": True, - "caption_projection_first_linear": False, - "caption_proj_input_norm": False, - "caption_projection_second_linear": False, -} - - -def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module: - """Select and create the appropriate feature extractor based on config. - Detection logic: - - V1: V2 config keys absent → projection lives in transformer - - V2: V2 config keys present with exact expected values → per-token RMS norm with dual aggregate embeds - - Anything else: NotImplementedError (config drift) - """ - gemma_text_config = GEMMA3_CONFIG_FOR_LTX.text_config - embedding_dim = gemma_text_config.hidden_size - num_layers = gemma_text_config.num_hidden_layers + 1 # +1 for the embedding layer - flat_dim = embedding_dim * num_layers - - overlapping_keys = transformer_config.keys() & _V2_EXPECTED_CONFIG.keys() - if not overlapping_keys: - aggregate_embed = torch.nn.Linear(flat_dim, embedding_dim, bias=False) - return FeatureExtractorV1(aggregate_embed=aggregate_embed, is_av=True) - - missing_keys = _V2_EXPECTED_CONFIG.keys() - overlapping_keys - if missing_keys: - raise NotImplementedError("Partial V2 config — missing keys: " + ", ".join(sorted(missing_keys))) - - unexpected_value_keys = {k for k in overlapping_keys if transformer_config[k] != _V2_EXPECTED_CONFIG[k]} - if unexpected_value_keys: - raise NotImplementedError("Unknown config: " + ", ".join(f"{k}={transformer_config[k]!r} (expected {_V2_EXPECTED_CONFIG[k]!r})" for k in unexpected_value_keys)) - - video_inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"] - audio_inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"] - return FeatureExtractorV2( - video_aggregate_embed=torch.nn.Linear(flat_dim, video_inner_dim, bias=True), - embedding_dim=embedding_dim, - audio_aggregate_embed=torch.nn.Linear(flat_dim, audio_inner_dim, bias=True), - ) - - -# --- Split SDOps: Gemma LLM keys vs Embeddings Processor keys --- - - -def _build_gemma_llm_key_ops(*, transformers_v5: bool) -> SDOps: - """Build the checkpoint-key remapping for the Gemma multimodal encoder. - The vision-tower mapping differs between transformers <5 and >=5 because - upstream PR https://github.com/huggingface/transformers/pull/39847 flattened - ``Gemma3ForConditionalGeneration.model.vision_tower.vision_model`` into - ``model.vision_tower``. Checkpoints continue to ship the legacy - ``vision_tower.vision_model.*`` prefix, so we strip the inner ``vision_model.`` - when targeting v5 and pass it through unchanged for v4. - """ - base = ( - SDOps("GEMMA_LLM_KEY_OPS") - # 1. Map language model layers (note the double .model prefix) - .with_matching(prefix="language_model.model.") - .with_replacement("language_model.model.", "model.model.language_model.") - # 2. Map the Vision Tower (version-dependent — see docstring) - .with_matching(prefix="vision_tower.") - ) - if transformers_v5: - base = base.with_replacement("vision_tower.vision_model.", "model.model.vision_tower.") - else: - base = base.with_replacement("vision_tower.", "model.model.vision_tower.") - return ( - base - # 3. Map the Multi-Modal Projector - .with_matching(prefix="multi_modal_projector.") - .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.") - # 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate()) - .with_kv_operation( - operation=lambda key, value: [ - KeyValueOperationResult(key, value), - KeyValueOperationResult("model.lm_head.weight", value), - ], - key_prefix="model.model.language_model.embed_tokens.weight", - ) - ) - - -GEMMA_LLM_KEY_OPS = _build_gemma_llm_key_ops(transformers_v5=_TRANSFORMERS_V5) - -EMBEDDINGS_PROCESSOR_KEY_OPS = ( - SDOps("EMBEDDINGS_PROCESSOR_KEY_OPS") - # 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor) - .with_matching(prefix="text_embedding_projection.aggregate_embed.") - .with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.") - # V2 dual aggregate embeds - .with_matching(prefix="text_embedding_projection.video_aggregate_embed.") - .with_replacement("text_embedding_projection.video_aggregate_embed.", "feature_extractor.video_aggregate_embed.") - .with_matching(prefix="text_embedding_projection.audio_aggregate_embed.") - .with_replacement("text_embedding_projection.audio_aggregate_embed.", "feature_extractor.audio_aggregate_embed.") - # 2. Map the connectors - .with_matching(prefix="model.diffusion_model.video_embeddings_connector.") - .with_replacement("model.diffusion_model.video_embeddings_connector.", "video_connector.") - .with_matching(prefix="model.diffusion_model.audio_embeddings_connector.") - .with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_connector.") -) - -VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = ( - SDOps("VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS") - # 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor) - .with_matching(prefix="text_embedding_projection.aggregate_embed.") - .with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.") - # V2 video aggregate embed - .with_matching(prefix="text_embedding_projection.video_aggregate_embed.") - .with_replacement("text_embedding_projection.video_aggregate_embed.", "feature_extractor.video_aggregate_embed.") - # 2. Map the connectors - .with_matching(prefix="model.diffusion_model.embeddings_connector.") - .with_replacement("model.diffusion_model.embeddings_connector.", "embeddings_processor.video_connector.") -) - - -def _resolve_local_base_freq(config: object) -> float: - rope_parameters = getattr(config, "rope_parameters", None) - if isinstance(rope_parameters, dict) and "sliding_attention" in rope_parameters: - sliding = rope_parameters["sliding_attention"] - if isinstance(sliding, dict) and "rope_theta" in sliding: - return float(sliding["rope_theta"]) - if hasattr(config, "rope_local_base_freq"): - return float(config.rope_local_base_freq) - raise AttributeError("Gemma text_config exposes neither rope_local_base_freq nor rope_parameters['sliding_attention']['rope_theta']") - - -def _resolve_full_rope_type(config: object) -> str: - rope_parameters = getattr(config, "rope_parameters", None) - if isinstance(rope_parameters, dict) and "full_attention" in rope_parameters: - full = rope_parameters["full_attention"] - if isinstance(full, dict) and "rope_type" in full: - return str(full["rope_type"]) - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is not None: - if isinstance(rope_scaling, dict): - if "rope_type" in rope_scaling: - return str(rope_scaling["rope_type"]) - elif hasattr(rope_scaling, "rope_type"): - return str(rope_scaling.rope_type) - raise AttributeError("Gemma text_config exposes neither rope_scaling.rope_type nor rope_parameters['full_attention']['rope_type']") - - -def _populate_rotary_v4(l_model: torch.nn.Module, config: object) -> None: - """transformers <5 layout: separate ``rotary_emb_local`` + ``rotary_emb``.""" - dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - base = _resolve_local_base_freq(config) - local_inv = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim)) - full_inv, _ = ROPE_INIT_FUNCTIONS[_resolve_full_rope_type(config)](config) - l_model.rotary_emb_local.register_buffer("inv_freq", local_inv) - l_model.rotary_emb.register_buffer("inv_freq", full_inv) - - -def _populate_rotary_v5(l_model: torch.nn.Module, config: object) -> None: - """transformers >=5 layout: single ``rotary_emb`` with per-layer-type buffers. - Mirrors ``Gemma3PreTrainedModel._init_weights`` for ``Gemma3RotaryEmbedding`` - so meta-built models reach the same numerical state as a from_pretrained load. - """ - rope_emb = l_model.rotary_emb - for layer_type in dict.fromkeys(config.layer_types): - rope_params = config.rope_parameters[layer_type] - if rope_params is None: - continue - rope_type = rope_params["rope_type"] - if rope_type == "default": - inv_freq, attn_scaling = rope_emb.compute_default_rope_parameters(config, layer_type=layer_type) - else: - inv_freq, attn_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, layer_type=layer_type) - rope_emb.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False) - rope_emb.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False) - setattr(rope_emb, f"{layer_type}_attention_scaling", attn_scaling) - - -def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder: - model = module.model - v_tower = model.model.vision_tower - v_model = v_tower.vision_model if hasattr(v_tower, "vision_model") else v_tower - l_model = model.model.language_model - - config = model.config.text_config - - if hasattr(l_model, "rotary_emb_local"): - _populate_rotary_v4(l_model, config) - else: - _populate_rotary_v5(l_model, config) - - positions_length = len(v_model.embeddings.position_ids[0]) - position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0) - v_model.embeddings.register_buffer("position_ids", position_ids) - embed_scale = torch.tensor(config.hidden_size**0.5, device="cpu") - l_model.embed_tokens.register_buffer("embed_scale", embed_scale) - - return module - - -GEMMA_MODEL_OPS = ModuleOps( - name="GemmaModel", - matcher=lambda module: hasattr(module, "model") and isinstance(module.model, Gemma3ForConditionalGeneration), - mutator=create_and_populate, -) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt b/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt deleted file mode 100644 index 0d677244f..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt +++ /dev/null @@ -1,30 +0,0 @@ -You are a Creative Assistant writing concise, action-focused image-to-video prompts. Given an image (first frame) and user Raw Input Prompt, generate a prompt to guide video generation from that image. - -#### Guidelines: -- Analyze the Image: Identify Subject, Setting, Elements, Style and Mood. -- Follow user Raw Input Prompt: Include all requested motion, actions, camera movements, audio, and details. If in conflict with the image, prioritize user request while maintaining visual consistency (describe transition from image to user's scene). -- Describe only changes from the image: Don't reiterate established visual details. Inaccurate descriptions may cause scene cuts. -- Active language: Use present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements. -- Chronological flow: Use temporal connectors ("as," "then," "while"). -- Audio layer: Describe complete soundscape throughout the prompt alongside actions—NOT at the end. Align audio intensity with action tempo. Include natural background audio, ambient sounds, effects, speech or music (when requested). Be specific (e.g., "soft footsteps on tile") not vague (e.g., "ambient sound"). -- Speech (only when requested): Provide exact words in quotes with character's visual/voice characteristics (e.g., "The tall man speaks in a low, gravelly voice"), language if not English and accent if relevant. If general conversation mentioned without text, generate contextual quoted dialogue. (i.e., "The man is talking" input -> the output should include exact spoken words, like: "The man is talking in an excited voice saying: 'You won't believe what I just saw!' His hands gesture expressively as he speaks, eyebrows raised with enthusiasm. The ambient sound of a quiet room underscores his animated speech.") -- Style: Include visual style at beginning: "Style: