From 53aaae8485340c07afcef02a1fbc659c49da1e97 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 3 Aug 2026 18:04:38 +0000 Subject: [PATCH 01/34] update train --- .../consistency/qwen_image_cm_cd_lora.yaml | 109 ++++++ .../consistency/qwen_image_cm_ct_lora.yaml | 104 +++++ .../qwen_image_mean_flow_ct_lora.yaml | 83 ++++ .../consistency/qwen_image_pcm_cd_lora.yaml | 103 +++++ .../consistency/qwen_image_scm_ct_lora.yaml | 81 ++++ .../consistency/qwen_image_tcm_ct_lora.yaml | 93 +++++ .../configs/train/flow/qwen_image_lora.yaml | 6 +- .../lightx2v_train/model_zoo/base.py | 122 +++++- .../lightx2v_train/model_zoo/qwen_image.py | 105 ++++- .../lightx2v_train/trainers/__init__.py | 2 + .../lightx2v_train/trainers/base.py | 1 + .../trainers/consistency/__init__.py | 24 ++ .../trainers/consistency/base.py | 188 +++++++++ .../lightx2v_train/trainers/consistency/cm.py | 302 +++++++++++++++ .../trainers/consistency/config.py | 365 ++++++++++++++++++ .../trainers/consistency/jvp.py | 30 ++ .../trainers/consistency/mean_flow.py | 287 ++++++++++++++ .../trainers/consistency/objective_factory.py | 29 ++ .../trainers/consistency/pcm.py | 237 ++++++++++++ .../trainers/consistency/scm.py | 223 +++++++++++ .../trainers/consistency/tcm.py | 155 ++++++++ .../trainers/consistency/trainer.py | 229 +++++++++++ .../lightx2v_train/trainers/flow.py | 4 + .../lightx2v_train/utils/registry.py | 2 + .../scripts/run_qwen_image_t2i_lora.sh | 1 + 25 files changed, 2856 insertions(+), 29 deletions(-) create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_cm_cd_lora.yaml create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_cm_ct_lora.yaml create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_mean_flow_ct_lora.yaml create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_scm_ct_lora.yaml create mode 100644 lightx2v_train/configs/train/consistency/qwen_image_tcm_ct_lora.yaml create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/__init__.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/base.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/cm.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/config.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/jvp.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/objective_factory.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/pcm.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/scm.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/tcm.py create mode 100644 lightx2v_train/lightx2v_train/trainers/consistency/trainer.py 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..6846da869 --- /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_native_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/flow/qwen_image_lora.yaml b/lightx2v_train/configs/train/flow/qwen_image_lora.yaml index a254124fb..e02555932 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 diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 0419e9eaa..71b94c1d2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -33,6 +33,22 @@ def dmd_latent_shape(self, batch_size, height, width): def denoiser_module(self): raise NotImplementedError(f"{self.__class__.__name__} must define denoiser_module().") + def configure_consistency_model(self, capabilities): + """Install optional architecture pieces required by an objective.""" + capabilities = frozenset(capabilities) + if capabilities: + names = ", ".join(sorted(capabilities)) + raise NotImplementedError(f"{self.__class__.__name__} does not support consistency capabilities: {names}.") + + def set_consistency_modules_trainable(self): + """Re-enable objective-specific modules after LoRA freezes the backbone.""" + + def consistency_auxiliary_parameter_names(self): + return () + + def predict_consistency_log_variance(self, time): + raise NotImplementedError(f"{self.__class__.__name__} does not implement a consistency log-variance head.") + def add_lora(self, rank, alpha, target_modules): lora_config = LoraConfig( r=rank, @@ -163,6 +179,34 @@ def denoise(self, denoiser_input, timesteps, condition): def postprocess_denoiser_output(self, prediction, denoiser_input): raise NotImplementedError + 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 distillation 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 @@ -205,7 +249,7 @@ def unload_lora_for_infer(self): 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) + peft_state_dict, auxiliary_state_dict = self._get_lora_and_auxiliary_state_dict_for_save(adapter_name=adapter_name) if not is_main_process(): return @@ -216,23 +260,38 @@ 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: + save_file( + auxiliary_state_dict, + os.path.join(output_dir, "consistency_auxiliary.safetensors"), + ) 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): 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(self.consistency_auxiliary_parameter_names()) + missing = auxiliary_names - state_dict.keys() + if missing: + raise RuntimeError(f"Consistency 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 +308,45 @@ 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 save_consistency_auxiliary_weights(self, save_dir): + names = set(self.consistency_auxiliary_parameter_names()) + if not names: + return + denoiser = self.denoiser_module() + if is_fsdp2_module(denoiser): + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + ignore_frozen_params=False, + strict=False, + ) + state_dict, _ = get_state_dict(denoiser, (), options=options) + else: + state_dict = denoiser.state_dict() + if not is_main_process(): + return + missing = names - state_dict.keys() + if missing: + raise RuntimeError(f"Consistency auxiliary parameters are missing from the model state: {sorted(missing)}") + auxiliary = {name: state_dict[name].detach().cpu().contiguous() for name in names if name in state_dict} + if auxiliary: + save_file(auxiliary, os.path.join(save_dir, "consistency_auxiliary.safetensors")) + + def load_consistency_auxiliary_weights(self, checkpoint_dir): + names = set(self.consistency_auxiliary_parameter_names()) + if not names: + return + path = os.path.join(checkpoint_dir, "consistency_auxiliary.safetensors") + if not os.path.exists(path): + raise RuntimeError(f"Consistency 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 consistency auxiliary keys in {path}: {missing}") + if unexpected: + logger.warning("Unexpected consistency auxiliary keys: {}", 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/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image.py index c16a26bbb..d3fea8380 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image.py @@ -1,8 +1,10 @@ +import copy from dataclasses import dataclass import torch from diffusers import AutoencoderKLQwenImage, QwenImagePipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor +from torch import nn from lightx2v_train.utils.registry import MODEL_REGISTER @@ -57,6 +59,62 @@ def load_transformer(self): def denoiser_module(self): return self.transformer + def denoiser_prediction_type(self): + # Qwen-Image follows x_t = (1 - t) * x_0 + t * noise and predicts + # the corresponding velocity noise - x_0. + return "velocity" + + def configure_consistency_model(self, capabilities): + capabilities = frozenset(capabilities) + unsupported = capabilities - {"endpoint_time", "log_variance"} + if unsupported: + names = ", ".join(sorted(unsupported)) + raise NotImplementedError(f"QwenImageModel does not support consistency capabilities: {names}.") + + transformer = self.transformer + if "log_variance" in capabilities and not hasattr(transformer, "logvar_linear"): + transformer.logvar_linear = nn.Linear(transformer.inner_dim, 1).to( + device=self.device, + dtype=self.running_dtype, + ) + + if "endpoint_time" in capabilities: + if not hasattr(transformer, "r_timestep_embedder"): + transformer.r_timestep_embedder = copy.deepcopy(transformer.time_text_embed.timestep_embedder) + self._consistency_endpoint_time = None + if "endpoint_time" in capabilities and not hasattr(self, "_consistency_endpoint_hook"): + self._consistency_endpoint_hook = transformer.time_text_embed.register_forward_hook(self._add_endpoint_time_embedding) + self._consistency_capabilities = getattr(self, "_consistency_capabilities", frozenset()) | capabilities + + def _add_endpoint_time_embedding(self, time_embedder, inputs, output): + endpoint_time = getattr(self, "_consistency_endpoint_time", None) + if endpoint_time is None: + return output + projected = time_embedder.time_proj(endpoint_time) + projected = projected.to(device=output.device, dtype=output.dtype) + endpoint_embedding = self.transformer.r_timestep_embedder(projected) + return output + endpoint_embedding.to(dtype=output.dtype) + + def set_consistency_modules_trainable(self): + transformer = self.transformer + for name in ("logvar_linear", "r_timestep_embedder"): + module = getattr(transformer, name, None) + if module is not None: + module.requires_grad_(True) + + def consistency_auxiliary_parameter_names(self): + prefixes = ("logvar_linear.", "r_timestep_embedder.") + return tuple(name for name, _ in self.transformer.named_parameters() if name.startswith(prefixes)) + + def predict_consistency_log_variance(self, time): + if not hasattr(self.transformer, "logvar_linear"): + raise RuntimeError("The Qwen consistency log-variance head has not been configured.") + time_embedder = self.transformer.time_text_embed + projected = time_embedder.time_proj(time) + dtype = self.transformer.logvar_linear.weight.dtype + embedding = time_embedder.timestep_embedder(projected.to(device=self.device, dtype=dtype)) + return self.transformer.logvar_linear(embedding) + def fsdp2_shard_plan(self, fsdp_config): reshard_config = fsdp_config["reshard_after_forward"] return [ @@ -107,16 +165,23 @@ def prepare_denoiser_input(self, noisy_latent, condition=None): width=w, ) - def denoise(self, denoiser_input, timestep_or_sigma, condition): - return self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, # timestep_or_sigma is in [0, 1] not [0, 1000] - guidance=None, - encoder_hidden_states_mask=condition["prompt_embed_mask"], - encoder_hidden_states=condition["prompt_embed"], - img_shapes=denoiser_input.img_shapes, - return_dict=False, - )[0] + def denoise(self, denoiser_input, timestep_or_sigma, condition, endpoint_time=None): + if endpoint_time is not None and not hasattr(self.transformer, "r_timestep_embedder"): + raise RuntimeError("endpoint_time was provided before Qwen endpoint conditioning was configured.") + previous_endpoint = getattr(self, "_consistency_endpoint_time", None) + self._consistency_endpoint_time = endpoint_time + try: + return self.transformer( + hidden_states=denoiser_input.hidden_states, + timestep=timestep_or_sigma, # timestep_or_sigma is in [0, 1] not [0, 1000] + guidance=None, + encoder_hidden_states_mask=condition["prompt_embed_mask"], + encoder_hidden_states=condition["prompt_embed"], + img_shapes=denoiser_input.img_shapes, + return_dict=False, + )[0] + finally: + self._consistency_endpoint_time = previous_endpoint def postprocess_denoiser_output(self, prediction, denoiser_input): return QwenImagePipeline._unpack_latents( @@ -165,9 +230,21 @@ def assemble_pipeline(self, scheduler=None): def get_pipeline_infer_kwargs(self, infer_config): # QwenImagePipeline uses `true_cfg_scale` instead of the standard `guidance_scale` - return { - "height": infer_config.get("height", 1024), - "width": infer_config.get("width", 1024), + enable_cfg = infer_config.get("enable_cfg", True) + kwargs = { + "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), - "true_cfg_scale": infer_config.get("cfg_guidance_scale", 4.0), + "true_cfg_scale": infer_config.get("cfg_guidance_scale", 4.0) if enable_cfg else 1.0, } + sigmas = infer_config.get("sigmas") + pcm_solver_steps = 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 not None: + from lightx2v_train.trainers.consistency.pcm import pcm_inference_sigmas + + sigmas = pcm_inference_sigmas(kwargs["num_inference_steps"], int(pcm_solver_steps)) + if sigmas is not None: + kwargs["sigmas"] = [float(value) for value in sigmas] + return kwargs diff --git a/lightx2v_train/lightx2v_train/trainers/__init__.py b/lightx2v_train/lightx2v_train/trainers/__init__.py index 77a6c0632..cdae33df5 100644 --- a/lightx2v_train/lightx2v_train/trainers/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/__init__.py @@ -4,6 +4,7 @@ _LAZY_EXPORTS = { "ARDmdTrainer": (".dmd.video_ar_trainer", "VideoArDmdTrainer"), + "ConsistencyTrainer": (".consistency.trainer", "ConsistencyTrainer"), "DmdTrainer": (".dmd.trainer", "DmdTrainer"), "DopsdTrainer": (".dopsd", "DopsdTrainer"), "FastWAMTrainer": (".fastwam", "FastWAMTrainer"), @@ -55,6 +56,7 @@ def __getattr__(name): __all__ = [ "build_trainer", "ARDmdTrainer", + "ConsistencyTrainer", "DmdTrainer", "FlowMatchingTrainer", "LTX2T2AVArDmdTrainer", diff --git a/lightx2v_train/lightx2v_train/trainers/base.py b/lightx2v_train/lightx2v_train/trainers/base.py index 1a670ede3..4234ec9ef 100644 --- a/lightx2v_train/lightx2v_train/trainers/base.py +++ b/lightx2v_train/lightx2v_train/trainers/base.py @@ -161,6 +161,7 @@ def _save_model_weights(self, model, save_dir): def _load_model_weights(self, model, save_dir): if self.train_type == "lora": model.load_lora_weights_for_resume(save_dir) + model.load_consistency_auxiliary_weights(save_dir) return model_state_path = os.path.join(save_dir, "model_state.pt") if not os.path.exists(model_state_path): diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py b/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py new file mode 100644 index 000000000..ad763f133 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py @@ -0,0 +1,24 @@ +from .base import ( + ConsistencyBatch, + ConsistencyObjective, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, + ReferenceModelSpec, +) +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER, build_consistency_objective + +__all__ = [ + "CONSISTENCY_OBJECTIVE_REGISTER", + "ConsistencyBatch", + "ConsistencyObjective", + "ConsistencyStepContext", + "DenoiserRequest", + "ModelDenoiser", + "ObjectiveOutput", + "RectifiedFlowPath", + "ReferenceModelSpec", + "build_consistency_objective", +] diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/base.py b/lightx2v_train/lightx2v_train/trainers/consistency/base.py new file mode 100644 index 000000000..413688b04 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/base.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional + +import torch +from torch import Tensor + +_PREDICTION_TYPE_ALIASES = { + "flow": "velocity", + "velocity": "velocity", + "x0": "x0", + "sample": "x0", + "noise": "noise", + "epsilon": "noise", +} + + +def canonical_prediction_type(prediction_type: str) -> str: + try: + return _PREDICTION_TYPE_ALIASES[prediction_type.lower()] + except (AttributeError, KeyError) as exc: + supported = ", ".join(sorted(_PREDICTION_TYPE_ALIASES)) + raise ValueError(f"Unsupported prediction type {prediction_type!r}; expected one of: {supported}.") from exc + + +def expand_time(time: Tensor, ndim: int) -> Tensor: + """Expand a scalar or batch time tensor over non-batch dimensions.""" + if time.ndim == 0: + time = time.reshape(1) + return time.reshape(time.shape[0], *([1] * (ndim - 1))) + + +class RectifiedFlowPath: + """Conversions for the straight path x_t=(1-t)x_0+t*noise.""" + + def interpolate(self, clean: Tensor, noise: Tensor, time: Tensor) -> Tensor: + time = expand_time(time, clean.ndim).to(device=clean.device, dtype=clean.dtype) + return (1.0 - time) * clean + time * noise + + def euler_step(self, sample: Tensor, velocity: Tensor, time: Tensor, next_time: Tensor) -> Tensor: + delta = expand_time(next_time - time, sample.ndim).to(device=sample.device, dtype=sample.dtype) + return (sample + delta * velocity.to(dtype=sample.dtype)).to(dtype=sample.dtype) + + def convert_prediction( + self, + sample: Tensor, + prediction: Tensor, + time: Tensor, + *, + source_type: str, + target_type: str, + ) -> Tensor: + source_type = canonical_prediction_type(source_type) + target_type = canonical_prediction_type(target_type) + if source_type == target_type: + return prediction + + time_expanded = expand_time(time, sample.ndim).to(device=sample.device, dtype=sample.dtype) + prediction = prediction.to(dtype=sample.dtype) + + if source_type == "velocity": + velocity = prediction + elif source_type == "x0": + denominator = torch.clamp(time_expanded, min=torch.finfo(sample.dtype).tiny) + velocity = (sample - prediction) / denominator + else: # noise + denominator = torch.clamp(1.0 - time_expanded, min=torch.finfo(sample.dtype).tiny) + velocity = (prediction - sample) / denominator + + if target_type == "velocity": + return velocity + if target_type == "x0": + return sample - time_expanded * velocity + return sample + (1.0 - time_expanded) * velocity + + +@dataclass(frozen=True) +class DenoiserRequest: + """Model-agnostic request issued by a consistency objective.""" + + sample: Tensor + time: Tensor + condition: Any + prediction_type: str = "velocity" + model_kwargs: Mapping[str, Any] = field(default_factory=dict) + + +class ModelDenoiser: + """Adapt a LightX2V model to the consistency objective interface.""" + + def __init__(self, model, path: RectifiedFlowPath): + self.model = model + self.path = path + + def predict(self, request: DenoiserRequest) -> Tensor: + raw_prediction = self.model.predict_denoiser_output( + request.sample, + request.time, + request.condition, + **request.model_kwargs, + ) + return self.path.convert_prediction( + request.sample, + raw_prediction, + request.time, + source_type=self.model.denoiser_prediction_type(), + target_type=request.prediction_type, + ) + + def predict_log_variance(self, time: Tensor) -> Tensor: + """Return the model-owned scalar log-variance head used by sCM.""" + predictor = getattr(self.model, "predict_consistency_log_variance", None) + if predictor is None: + raise NotImplementedError(f"{self.model.__class__.__name__} does not provide a consistency log-variance head.") + return predictor(time) + + +@dataclass(frozen=True) +class ConsistencyBatch: + clean: Tensor + condition: Any + negative_condition: Any = None + + +@dataclass(frozen=True) +class ConsistencyStepContext: + iteration: int + global_batch_size: int + + +@dataclass +class ObjectiveOutput: + loss: Tensor + metrics: Mapping[str, Tensor | float] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ReferenceModelSpec: + """Description of an algorithm-owned frozen model. + + ``checkpoint`` is a LightX2V training checkpoint directory. Model + overrides are read from ``model.`` so a future objective can change + the reference architecture without teaching the trainer about it. + """ + + role: str + checkpoint: str + training_mode: bool = False + + +class ConsistencyObjective(ABC): + """Extension point implemented by CM, sCM, TCM, PCM, MeanFlow, and others.""" + + algorithm_name = "base" + requires_teacher = False + requires_negative_condition = False + negative_prompt: Optional[str] = None + model_capabilities: frozenset[str] = frozenset() + + @property + def reference_model_specs(self) -> tuple[ReferenceModelSpec, ...]: + return () + + @property + def student_initialization_checkpoint(self) -> Optional[str]: + return None + + @abstractmethod + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + """Sample all stochastic state before sequence-parallel broadcast.""" + + @abstractmethod + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + """Compute a differentiable scalar loss and detached metrics.""" diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/cm.py b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py new file mode 100644 index 000000000..fd57e8a37 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py @@ -0,0 +1,302 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +from torch import Tensor + +from .base import ( + ConsistencyBatch, + ConsistencyObjective, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, + expand_time, +) +from .config import CMConfig, CMLossConfig, CMTimePairConfig +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER + + +@dataclass(frozen=True) +class CMTimePair: + t: Tensor + r: Tensor + ratio: float + stage: int + + +@dataclass(frozen=True) +class CMLossTerms: + weighted: Tensor + unweighted: Tensor + prediction: Tensor + target: Tensor + + +class CMTimePairSampler: + """Build the adjacent time pair used by the CM curriculum. + + The default ``ect`` mapping and ratio curriculum match FastGen's CM + implementation. ``linear`` is also available for rectified-flow + experiments without changing the objective itself. + """ + + def __init__(self, config: CMTimePairConfig): + self.config = config + + def sample(self, t: Tensor, context: ConsistencyStepContext) -> CMTimePair: + stage = int((context.iteration * context.global_batch_size) // (self.config.kimg_per_stage * 1000.0)) + ratio = min(1.0 - 1.0 / self.config.q ** (stage + 1), self.config.ratio_limit) + + # Pair construction remains in fp32 even when the model runs in bf16; + # otherwise late curriculum stages can round t and r to the same time. + t = t.float() + if self.config.mapping in {"ect", "sigmoid"}: + r = t - t * (1.0 - ratio) * (1.0 + 8.0 * torch.sigmoid(-t)) + else: + r = ratio * t + + r = torch.clamp(r, min=self.config.min_r) + # Keep a strictly positive interval for the loss weighting. The + # boundary r=0 remains valid and is handled exactly by the objective. + largest_safe_r = torch.clamp(t - self.config.safety_epsilon, min=0.0) + r = torch.minimum(r, largest_safe_r) + return CMTimePair(t=t, r=r, ratio=ratio, stage=stage) + + +class CMLoss: + """Vector-distance CM loss with configurable interval weighting.""" + + _WEIGHT_ALIASES = { + "default": "inverse_delta", + "sqrt": "inverse_sqrt_delta", + "one": "none", + } + + def __init__(self, config: CMLossConfig): + self.config = config + self.dtype = torch.float64 if config.computation_dtype == "float64" else torch.float32 + + def __call__(self, prediction: Tensor, target: Tensor, t: Tensor, r: Tensor): + difference = prediction.to(self.dtype) - target.to(self.dtype) + squared_distance = difference.flatten(1).square() + if self.config.normalize_by_numel: + squared_distance = squared_distance.mean(dim=1) + else: + squared_distance = squared_distance.sum(dim=1) + + if self.config.distance == "squared_l2": + unweighted = squared_distance + else: + l2_distance = torch.sqrt(squared_distance) + if self.config.distance == "pseudo_huber": + constant = self.config.huber_constant + unweighted = torch.sqrt(l2_distance.square() + constant**2) - constant + else: + unweighted = l2_distance + + delta = (t.to(self.dtype) - r.to(self.dtype)).clamp_min(self.config.min_denominator) + weighting = self._WEIGHT_ALIASES.get(self.config.weighting, self.config.weighting) + if weighting == "inverse_delta": + weighted = unweighted / delta + elif weighting == "inverse_sqrt_delta": + weighted = unweighted / torch.sqrt(delta) + elif weighting in {"c_out", "c_out_sq"}: + # EDM c_out after matching the RF signal/noise ratio + # sigma_edm=t/(1-t), written in a form that stays finite at t=1. + time = t.to(self.dtype) + sigma_data = self.config.sigma_data + c_out = time * sigma_data / torch.sqrt(time.square() + sigma_data**2 * (1.0 - time).square()).clamp_min(self.config.min_denominator) + denominator = c_out.square() if weighting == "c_out_sq" else c_out + weighted = unweighted / denominator.clamp_min(self.config.min_denominator) + elif weighting == "sigma_sq": + weighted = unweighted / t.to(self.dtype).square().clamp_min(self.config.min_denominator) + else: + weighted = unweighted + return weighted, unweighted + + +def classifier_free_guidance( + conditional: Tensor, + unconditional: Tensor, + scale: float, + norm: str, +) -> Tensor: + guided = unconditional + scale * (conditional - unconditional) + if norm == "none": + return guided + if norm == "layer_norm": + conditional_norm = torch.linalg.vector_norm(conditional, dim=-1, keepdim=True) + guided_norm = torch.linalg.vector_norm(guided, dim=-1, keepdim=True) + return guided * (conditional_norm / guided_norm.clamp_min(1e-12)) + conditional_norm = torch.linalg.vector_norm(conditional) + guided_norm = torch.linalg.vector_norm(guided) + scale_correction = torch.clamp(conditional_norm / guided_norm.clamp_min(1e-12), max=1.0) + return guided * scale_correction + + +@CONSISTENCY_OBJECTIVE_REGISTER("cm") +class CMObjective(ConsistencyObjective): + """Consistency Training (CT) and Consistency Distillation (CD).""" + + algorithm_name = "cm" + + def __init__(self, config, path: RectifiedFlowPath): + self.config = CMConfig.from_mapping(config) + self.path = path + self.time_pair_sampler = CMTimePairSampler(self.config.time_pair) + self.loss_fn = CMLoss(self.config.loss) + self.requires_teacher = self.config.mode == "cd" + guidance_scale = self.config.teacher.guidance_scale + self.requires_negative_condition = self.requires_teacher and guidance_scale is not None and guidance_scale != 1.0 + self.negative_prompt = self.config.teacher.negative_prompt + + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + batch_size = clean.shape[0] + latent_hw = (clean.shape[-2], clean.shape[-1]) + sampled_t = scheduler.sample_timestep_or_sigma(batch_size, latent_hw=latent_hw).to(clean.device) + pair = self.time_pair_sampler.sample(sampled_t, context) + return { + "noise": torch.randn_like(clean), + "t": pair.t, + "r": pair.r, + "ratio": torch.tensor(pair.ratio, device=clean.device, dtype=torch.float32), + "stage": torch.tensor(pair.stage, device=clean.device, dtype=torch.float32), + } + + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + del references + terms = self.compute_loss_terms(batch, training_state, student, teacher) + t = training_state["t"] + r = training_state["r"] + delta = t.float() - r.float() + return ObjectiveOutput( + loss=terms.weighted.mean(), + metrics={ + "cm_unweighted": terms.unweighted.detach().mean(), + "cm_t": t.detach().float().mean(), + "cm_r": r.detach().float().mean(), + "cm_delta": delta.detach().mean(), + "cm_ratio": training_state["ratio"].detach(), + "cm_stage": training_state["stage"].detach(), + }, + ) + + def compute_loss_terms( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + ) -> CMLossTerms: + clean = batch.clean + noise = training_state["noise"] + t = training_state["t"] + r = training_state["r"] + noisy_t = self.path.interpolate(clean, noise, t) + + if self.config.mode == "ct": + noisy_r = self.path.interpolate(clean, noise, r) + else: + if teacher is None: + raise RuntimeError("CM consistency distillation requires a frozen teacher denoiser.") + noisy_r = self._distill_to_r( + noisy_t, + t, + r, + batch.condition, + batch.negative_condition, + teacher, + ) + + # The first forward is forked so the target forward sees exactly the + # same dropout RNG state. Exiting fork_rng restores that state before + # evaluating the stop-gradient target branch. + rng_devices = [clean.device] if clean.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices): + prediction = student.predict( + DenoiserRequest( + sample=noisy_t, + time=t, + condition=batch.condition, + prediction_type="x0", + ) + ) + with torch.no_grad(): + target_candidate = student.predict( + DenoiserRequest( + sample=noisy_r, + time=r, + condition=batch.condition, + prediction_type="x0", + ) + ) + + # D(x, 0)=x is the CM boundary condition. For data-paired CT/CD, + # FastGen anchors an r=0 target to the clean sample explicitly. + positive_r = expand_time(r > 0, clean.ndim) + target_candidate = torch.nan_to_num(target_candidate) + target = torch.where(positive_r, target_candidate, clean) + + weighted, unweighted = self.loss_fn(prediction, target, t, r) + return CMLossTerms( + weighted=weighted, + unweighted=unweighted, + prediction=prediction, + target=target, + ) + + @torch.no_grad() + def _distill_to_r( + self, + noisy_t: Tensor, + t: Tensor, + r: Tensor, + condition, + negative_condition, + teacher: ModelDenoiser, + ) -> Tensor: + conditional_velocity = teacher.predict( + DenoiserRequest( + sample=noisy_t, + time=t, + condition=condition, + prediction_type="velocity", + ) + ) + guidance_scale = self.config.teacher.guidance_scale + if guidance_scale is None or guidance_scale == 1.0: + velocity = conditional_velocity + else: + if negative_condition is None: + raise RuntimeError("CM teacher CFG requires a negative condition.") + unconditional_velocity = teacher.predict( + DenoiserRequest( + sample=noisy_t, + time=t, + condition=negative_condition, + prediction_type="velocity", + ) + ) + velocity = classifier_free_guidance( + conditional_velocity, + unconditional_velocity, + guidance_scale, + self.config.teacher.cfg_norm, + ) + return self.path.euler_step(noisy_t, velocity, t, r) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/config.py b/lightx2v_train/lightx2v_train/trainers/consistency/config.py new file mode 100644 index 000000000..5b6474b92 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/config.py @@ -0,0 +1,365 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Optional + + +def _mapping(value, name): + if value is None: + return {} + if not isinstance(value, Mapping): + raise ValueError(f"{name} must be a mapping.") + return value + + +@dataclass(frozen=True) +class CMTimePairConfig: + mapping: str = "ect" + q: float = 2.0 + ratio_limit: float = 0.999 + kimg_per_stage: float = 12500.0 + min_r: float = 0.0 + safety_epsilon: float = 1e-6 + + @classmethod + def from_mapping(cls, value): + value = _mapping(value, "training.consistency.time_pair") + result = cls( + mapping=str(value.get("mapping", "ect")).lower(), + q=float(value.get("q", 2.0)), + ratio_limit=float(value.get("ratio_limit", 0.999)), + kimg_per_stage=float(value.get("kimg_per_stage", 12500.0)), + min_r=float(value.get("min_r", 0.0)), + safety_epsilon=float(value.get("safety_epsilon", 1e-6)), + ) + if result.mapping not in {"ect", "sigmoid", "linear"}: + raise ValueError("training.consistency.time_pair.mapping must be 'ect', 'sigmoid', or 'linear'.") + if result.q <= 1.0: + raise ValueError("training.consistency.time_pair.q must be greater than 1.") + if not 0.0 <= result.ratio_limit < 1.0: + raise ValueError("training.consistency.time_pair.ratio_limit must be in [0, 1).") + if result.kimg_per_stage <= 0.0: + raise ValueError("training.consistency.time_pair.kimg_per_stage must be positive.") + if result.min_r < 0.0: + raise ValueError("training.consistency.time_pair.min_r must be non-negative.") + if result.safety_epsilon <= 0.0: + raise ValueError("training.consistency.time_pair.safety_epsilon must be positive.") + return result + + +@dataclass(frozen=True) +class CMLossConfig: + distance: str = "pseudo_huber" + huber_constant: float = 1e-8 + weighting: str = "inverse_delta" + normalize_by_numel: bool = False + computation_dtype: str = "float32" + min_denominator: float = 1e-12 + sigma_data: float = 0.5 + + @classmethod + def from_mapping(cls, value): + value = _mapping(value, "training.consistency.loss") + result = cls( + distance=str(value.get("distance", "pseudo_huber")).lower(), + huber_constant=float(value.get("huber_constant", 1e-8)), + weighting=str(value.get("weighting", "inverse_delta")).lower(), + normalize_by_numel=bool(value.get("normalize_by_numel", False)), + computation_dtype=str(value.get("computation_dtype", "float32")).lower(), + min_denominator=float(value.get("min_denominator", 1e-12)), + sigma_data=float(value.get("sigma_data", 0.5)), + ) + if result.distance not in {"pseudo_huber", "l2", "squared_l2"}: + raise ValueError("training.consistency.loss.distance must be 'pseudo_huber', 'l2', or 'squared_l2'.") + if result.huber_constant < 0.0: + raise ValueError("training.consistency.loss.huber_constant must be non-negative.") + if result.weighting not in { + "inverse_delta", + "inverse_sqrt_delta", + "none", + "default", + "sqrt", + "one", + "c_out", + "c_out_sq", + "sigma_sq", + }: + raise ValueError("Unsupported training.consistency.loss.weighting.") + if result.computation_dtype not in {"float32", "float64"}: + raise ValueError("training.consistency.loss.computation_dtype must be 'float32' or 'float64'.") + if result.min_denominator <= 0.0: + raise ValueError("training.consistency.loss.min_denominator must be positive.") + if result.sigma_data <= 0.0: + raise ValueError("training.consistency.loss.sigma_data must be positive.") + return result + + +@dataclass(frozen=True) +class CMTeacherConfig: + guidance_scale: Optional[float] = None + negative_prompt: str = " " + cfg_norm: str = "none" + + @classmethod + def from_mapping(cls, value): + value = _mapping(value, "training.consistency.teacher") + raw_scale = value.get("guidance_scale") + result = cls( + guidance_scale=None if raw_scale is None else float(raw_scale), + negative_prompt=str(value.get("negative_prompt", " ")), + cfg_norm=str(value.get("cfg_norm", "none")).lower(), + ) + if result.guidance_scale is not None and result.guidance_scale < 0.0: + raise ValueError("training.consistency.teacher.guidance_scale must be non-negative.") + if result.cfg_norm not in {"none", "layer_norm", "scalar"}: + raise ValueError("training.consistency.teacher.cfg_norm must be 'none', 'layer_norm', or 'scalar'.") + return result + + +@dataclass(frozen=True) +class CMConfig: + mode: str + time_pair: CMTimePairConfig + loss: CMLossConfig + teacher: CMTeacherConfig + + @classmethod + def from_mapping(cls, config): + consistency = _mapping(config["training"].get("consistency"), "training.consistency") + mode = str(consistency.get("mode", "ct")).lower() + if mode not in {"ct", "cd"}: + raise ValueError("training.consistency.mode must be 'ct' or 'cd'.") + return cls( + mode=mode, + time_pair=CMTimePairConfig.from_mapping(consistency.get("time_pair")), + loss=CMLossConfig.from_mapping(consistency.get("loss")), + teacher=CMTeacherConfig.from_mapping(consistency.get("teacher")), + ) + + +@dataclass(frozen=True) +class JVPConfig: + method: str = "finite_difference" + epsilon: float = 1e-3 + + @classmethod + def from_mapping(cls, value, *, default_epsilon): + value = _mapping(value, "training.consistency.jvp") + method = str(value.get("method", "finite_difference")).lower() + aliases = {"fd": "finite_difference", "finite-difference": "finite_difference"} + result = cls( + method=aliases.get(method, method), + epsilon=float(value.get("epsilon", default_epsilon)), + ) + if result.method not in {"finite_difference", "exact"}: + raise ValueError("training.consistency.jvp.method must be 'finite_difference' or 'exact'.") + if result.epsilon <= 0.0: + raise ValueError("training.consistency.jvp.epsilon must be positive.") + return result + + +@dataclass(frozen=True) +class SCMConfig: + mode: str + sigma_data: float + tangent_warmup_steps: int + tangent_warmup_constant: float + prior_weighting: bool + spatially_normalized_tangent: bool + normalize_by_numel: bool + min_denominator: float + jvp: JVPConfig + teacher: CMTeacherConfig + + @classmethod + def from_mapping(cls, config): + consistency = _mapping(config["training"].get("consistency"), "training.consistency") + loss = _mapping(consistency.get("loss"), "training.consistency.loss") + mode = str(consistency.get("mode", "ct")).lower() + if mode not in {"ct", "cd"}: + raise ValueError("training.consistency.mode must be 'ct' or 'cd'.") + result = cls( + mode=mode, + sigma_data=float(consistency.get("sigma_data", 0.5)), + tangent_warmup_steps=int(loss.get("tangent_warmup_steps", 10000)), + tangent_warmup_constant=float(loss.get("tangent_warmup_constant", 0.1)), + prior_weighting=bool(loss.get("prior_weighting", True)), + spatially_normalized_tangent=bool(loss.get("spatially_normalized_tangent", True)), + normalize_by_numel=bool(loss.get("normalize_by_numel", True)), + min_denominator=float(loss.get("min_denominator", 1e-12)), + jvp=JVPConfig.from_mapping(consistency.get("jvp"), default_epsilon=1e-3), + teacher=CMTeacherConfig.from_mapping(consistency.get("teacher")), + ) + if result.sigma_data <= 0.0: + raise ValueError("training.consistency.sigma_data must be positive.") + if result.tangent_warmup_steps < 0: + raise ValueError("tangent_warmup_steps must be non-negative.") + if result.tangent_warmup_constant <= 0.0 or result.min_denominator <= 0.0: + raise ValueError("sCM normalization constants must be positive.") + return result + + +@dataclass(frozen=True) +class TCMConfig: + cm: CMConfig + transition_time: float + boundary_probability: float + boundary_weight: float + stage1_checkpoint: str + + @classmethod + def from_mapping(cls, config): + consistency = _mapping(config["training"].get("consistency"), "training.consistency") + stage1_checkpoint = str(consistency.get("stage1_checkpoint", "")).strip() + result = cls( + cm=CMConfig.from_mapping(config), + transition_time=float(consistency.get("transition_time", 0.5)), + boundary_probability=float(consistency.get("boundary_probability", 0.25)), + boundary_weight=float(consistency.get("boundary_weight", 0.1)), + stage1_checkpoint=stage1_checkpoint, + ) + if not 0.0 < result.transition_time <= 1.0: + raise ValueError("training.consistency.transition_time must be in (0, 1].") + if not 0.0 <= result.boundary_probability <= 1.0: + raise ValueError("training.consistency.boundary_probability must be in [0, 1].") + if result.boundary_weight < 0.0: + raise ValueError("training.consistency.boundary_weight must be non-negative.") + if not result.stage1_checkpoint: + raise ValueError("TCM requires training.consistency.stage1_checkpoint.") + return result + + +@dataclass(frozen=True) +class PCMSolverConfig: + num_solver_steps: int + num_phases: int + boundary_time: Optional[float] + + @classmethod + def from_mapping(cls, value): + value = _mapping(value, "training.consistency.solver") + raw_boundary = value.get("boundary_time") + result = cls( + num_solver_steps=int(value.get("num_solver_steps", 50)), + num_phases=int(value.get("num_phases", 4)), + boundary_time=None if raw_boundary is None else float(raw_boundary), + ) + if result.num_solver_steps <= 0: + raise ValueError("PCM solver.num_solver_steps must be positive.") + if not 1 <= result.num_phases <= result.num_solver_steps: + raise ValueError("PCM solver.num_phases must be in [1, num_solver_steps].") + if result.boundary_time is not None and not 0.0 <= result.boundary_time < 1.0: + raise ValueError("PCM solver.boundary_time must be in [0, 1).") + return result + + +@dataclass(frozen=True) +class PCMLossConfig: + distance: str + huber_constant: float + computation_dtype: str + + @classmethod + def from_mapping(cls, value): + value = _mapping(value, "training.consistency.loss") + distance = str(value.get("distance", "pseudo_huber")).lower() + aliases = {"huber": "pseudo_huber", "mse": "l2", "squared_l2": "l2"} + result = cls( + distance=aliases.get(distance, distance), + huber_constant=float(value.get("huber_constant", 1e-3)), + computation_dtype=str(value.get("computation_dtype", "float32")).lower(), + ) + if result.distance not in {"pseudo_huber", "l2"}: + raise ValueError("PCM loss.distance must be 'pseudo_huber' or 'l2'.") + if result.huber_constant < 0.0: + raise ValueError("PCM loss.huber_constant must be non-negative.") + if result.computation_dtype not in {"float32", "float64"}: + raise ValueError("PCM loss.computation_dtype must be 'float32' or 'float64'.") + return result + + +@dataclass(frozen=True) +class PCMConfig: + mode: str + solver: PCMSolverConfig + loss: PCMLossConfig + teacher: CMTeacherConfig + + @classmethod + def from_mapping(cls, config): + consistency = _mapping(config["training"].get("consistency"), "training.consistency") + mode = str(consistency.get("mode", "cd")).lower() + if mode != "cd": + raise ValueError( + "The published PCM algorithm is distillation-only; " + "set training.consistency.mode to 'cd'." + ) + return cls( + mode=mode, + solver=PCMSolverConfig.from_mapping(consistency.get("solver")), + loss=PCMLossConfig.from_mapping(consistency.get("loss")), + teacher=CMTeacherConfig.from_mapping(consistency.get("teacher")), + ) + + +@dataclass(frozen=True) +class MeanFlowConfig: + mode: str + random_endpoint_probability: float + loss_type: str + norm_method: str + norm_constant: float + tangent_warmup_steps: int + spatially_normalized_tangent: bool + min_denominator: float + condition_dropout_probability: Optional[float] + guidance_scale: Optional[float] + guidance_mixture_ratio: Optional[float] + guidance_time_start: float + guidance_time_end: float + jvp: JVPConfig + teacher: CMTeacherConfig + + @classmethod + def from_mapping(cls, config): + consistency = _mapping(config["training"].get("consistency"), "training.consistency") + loss = _mapping(consistency.get("loss"), "training.consistency.loss") + sampling = _mapping(consistency.get("sampling"), "training.consistency.sampling") + guidance = _mapping(consistency.get("guidance"), "training.consistency.guidance") + teacher = _mapping(consistency.get("teacher"), "training.consistency.teacher") + mode = str(consistency.get("mode", "ct")).lower() + raw_dropout = guidance.get("condition_dropout_probability") + raw_scale = guidance.get("scale", teacher.get("guidance_scale")) + raw_mixture = guidance.get("mixture_ratio") + result = cls( + mode=mode, + random_endpoint_probability=float(sampling.get("random_endpoint_probability", 0.0)), + loss_type=str(loss.get("type", "opt_grad")).lower(), + norm_method=str(loss.get("norm_method", "poly_1.0")).lower(), + norm_constant=float(loss.get("norm_constant", 0.1)), + tangent_warmup_steps=int(loss.get("tangent_warmup_steps", 0)), + spatially_normalized_tangent=bool(loss.get("spatially_normalized_tangent", False)), + min_denominator=float(loss.get("min_denominator", 1e-12)), + condition_dropout_probability=None if raw_dropout is None else float(raw_dropout), + guidance_scale=None if raw_scale is None else float(raw_scale), + guidance_mixture_ratio=None if raw_mixture is None else float(raw_mixture), + guidance_time_start=float(guidance.get("time_start", 0.0)), + guidance_time_end=float(guidance.get("time_end", 1.0)), + jvp=JVPConfig.from_mapping(consistency.get("jvp"), default_epsilon=1e-4), + teacher=CMTeacherConfig.from_mapping(consistency.get("teacher")), + ) + if result.mode not in {"ct", "cd"}: + raise ValueError("training.consistency.mode must be 'ct' or 'cd'.") + if not 0.0 <= result.random_endpoint_probability <= 1.0: + raise ValueError("random_endpoint_probability must be in [0, 1].") + if result.loss_type not in {"l2", "opt_grad"}: + raise ValueError("MeanFlow loss.type must be 'l2' or 'opt_grad'.") + if result.tangent_warmup_steps < 0 or result.norm_constant <= 0.0 or result.min_denominator <= 0.0: + raise ValueError("MeanFlow normalization values are invalid.") + if result.condition_dropout_probability is not None and not 0.0 <= result.condition_dropout_probability <= 1.0: + raise ValueError("condition_dropout_probability must be in [0, 1].") + if result.guidance_scale is not None and result.guidance_scale < 0.0: + raise ValueError("MeanFlow guidance scale must be non-negative.") + if result.guidance_time_start > result.guidance_time_end: + raise ValueError("guidance.time_start cannot be greater than guidance.time_end.") + return result diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/jvp.py b/lightx2v_train/lightx2v_train/trainers/consistency/jvp.py new file mode 100644 index 000000000..a8159452a --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/jvp.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import contextlib + +import torch + + +@contextlib.contextmanager +def math_attention_for_forward_ad(device_type: str): + """Use the SDPA backend that supports forward-mode automatic differentiation.""" + if device_type != "cuda": + yield + return + + backend = torch.backends.cuda + flash = backend.flash_sdp_enabled() + memory_efficient = backend.mem_efficient_sdp_enabled() + cudnn = backend.cudnn_sdp_enabled() + math = backend.math_sdp_enabled() + backend.enable_flash_sdp(False) + backend.enable_mem_efficient_sdp(False) + backend.enable_cudnn_sdp(False) + backend.enable_math_sdp(True) + try: + yield + finally: + backend.enable_flash_sdp(flash) + backend.enable_mem_efficient_sdp(memory_efficient) + backend.enable_cudnn_sdp(cudnn) + backend.enable_math_sdp(math) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py new file mode 100644 index 000000000..ad5c36418 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +from typing import Any, Mapping, Optional + +import torch +from torch import Tensor + +from .base import ( + ConsistencyBatch, + ConsistencyObjective, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, + expand_time, +) +from .config import MeanFlowConfig +from .jvp import math_attention_for_forward_ad +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER + + +def _where_condition(mask: Tensor, positive: Any, negative: Any): + if torch.is_tensor(positive): + return torch.where(expand_time(mask, positive.ndim), positive, negative) + if isinstance(positive, Mapping): + return {key: _where_condition(mask, value, negative[key]) for key, value in positive.items()} + raise TypeError(f"MeanFlow condition dropout does not support {type(positive).__name__}.") + + +@CONSISTENCY_OBJECTIVE_REGISTER("mean_flow") +@CONSISTENCY_OBJECTIVE_REGISTER("meanflow") +class MeanFlowObjective(ConsistencyObjective): + """MeanFlow objective with endpoint conditioning and JVP/finite differences.""" + + algorithm_name = "mean_flow" + model_capabilities = frozenset({"endpoint_time"}) + + def __init__(self, config, path: RectifiedFlowPath): + self.config = MeanFlowConfig.from_mapping(config) + self.path = path + self.requires_teacher = self.config.mode == "cd" + uses_training_guidance = any( + value is not None + for value in ( + self.config.condition_dropout_probability, + self.config.guidance_scale, + self.config.guidance_mixture_ratio, + ) + ) + self.requires_negative_condition = uses_training_guidance + self.negative_prompt = self.config.teacher.negative_prompt + + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + latent_hw = (clean.shape[-2], clean.shape[-1]) + first = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() + second = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() + t = torch.maximum(first, second) + random_r = torch.minimum(first, second) + use_random_r = torch.rand(clean.shape[0], device=clean.device) < self.config.random_endpoint_probability + r = torch.where(use_random_r, random_r, t) + return { + "noise": torch.randn_like(clean), + "t": t, + "r": r, + "random_endpoint_mask": use_random_r, + "iteration": torch.tensor(context.iteration, device=clean.device, dtype=torch.int64), + } + + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + del references + clean = batch.clean + noise = training_state["noise"] + t = training_state["t"].float() + r = training_state["r"].float() + x_t = self.path.interpolate(clean, noise, t) + condition, dxt_dt = self._target_velocity(batch, x_t, noise, t, student, teacher) + + velocity_jvp = self._jvp(x_t, t, r, dxt_dt, condition, student) + velocity = student.predict( + DenoiserRequest( + sample=x_t, + time=t, + condition=condition, + prediction_type="velocity", + model_kwargs={"endpoint_time": r}, + ) + ) + loss, tangent, weight, warmup = self._loss( + velocity, + velocity_jvp, + dxt_dt, + t, + r, + int(training_state["iteration"].item()), + ) + flow_matching_error = (velocity.float() - (noise.float() - clean.float())).square().flatten(1).mean(dim=1) + return ObjectiveOutput( + loss=loss.mean(), + metrics={ + "mean_flow_loss": loss.detach().mean(), + "mean_flow_velocity_mse": flow_matching_error.detach().mean(), + "mean_flow_jvp_norm": torch.linalg.vector_norm(velocity_jvp.flatten(1), dim=1).detach().mean(), + "mean_flow_tangent_norm": torch.linalg.vector_norm(tangent.flatten(1), dim=1).detach().mean(), + "mean_flow_weight": weight.detach().mean(), + "mean_flow_t": t.detach().mean(), + "mean_flow_r": r.detach().mean(), + "mean_flow_random_endpoint_fraction": training_state["random_endpoint_mask"].float().mean(), + "mean_flow_warmup": warmup, + }, + ) + + @torch.no_grad() + def _target_velocity( + self, + batch: ConsistencyBatch, + x_t: Tensor, + noise: Tensor, + t: Tensor, + student: ModelDenoiser, + teacher: Optional[ModelDenoiser], + ): + if self.config.mode == "cd": + if teacher is None: + raise RuntimeError("MeanFlow consistency distillation requires a frozen teacher denoiser.") + velocity = teacher.predict(DenoiserRequest(x_t, t, batch.condition, prediction_type="velocity")) + if self.config.guidance_scale is not None: + if batch.negative_condition is None: + raise RuntimeError("MeanFlow teacher CFG requires a negative condition.") + negative = teacher.predict(DenoiserRequest(x_t, t, batch.negative_condition, prediction_type="velocity")) + scale = self._time_limited_value(t, self.config.guidance_scale, outside=1.0) + velocity = negative + expand_time(scale, velocity.ndim) * (velocity - negative) + return batch.condition, velocity + + velocity = noise - batch.clean + if self.config.guidance_scale is None and self.config.guidance_mixture_ratio is None: + return batch.condition, velocity + if batch.negative_condition is None: + raise RuntimeError("MeanFlow training guidance requires a negative condition.") + + negative = student.predict( + DenoiserRequest( + x_t, + t, + batch.negative_condition, + prediction_type="velocity", + model_kwargs={"endpoint_time": t}, + ) + ) + scale_value = 1.0 if self.config.guidance_scale is None else self.config.guidance_scale + scale = self._time_limited_value(t, scale_value, outside=1.0) + if self.config.guidance_mixture_ratio is None: + guided = negative + expand_time(scale, velocity.ndim) * (velocity - negative) + else: + conditional = student.predict( + DenoiserRequest( + x_t, + t, + batch.condition, + prediction_type="velocity", + model_kwargs={"endpoint_time": t}, + ) + ) + mixture = self._time_limited_value(t, self.config.guidance_mixture_ratio, outside=0.0) + guided = expand_time(scale, velocity.ndim) * velocity + expand_time(1.0 - scale - mixture, velocity.ndim) * negative + expand_time(mixture, velocity.ndim) * conditional + + dropout = self.config.condition_dropout_probability + if dropout is None: + return batch.condition, guided + keep_condition = torch.rand(t.shape[0], device=t.device) >= dropout + mixed_condition = _where_condition(keep_condition, batch.condition, batch.negative_condition) + mixed_velocity = torch.where(expand_time(keep_condition, velocity.ndim), guided, velocity) + return mixed_condition, mixed_velocity + + def _time_limited_value(self, t: Tensor, value: float, *, outside: float) -> Tensor: + active = (t >= self.config.guidance_time_start) & (t <= self.config.guidance_time_end) + return torch.where(active, torch.full_like(t, value), torch.full_like(t, outside)) + + @torch.no_grad() + def _jvp( + self, + x_t: Tensor, + t: Tensor, + r: Tensor, + dxt_dt: Tensor, + condition, + student: ModelDenoiser, + ) -> Tensor: + def model_fn(sample, time, endpoint): + return student.predict( + DenoiserRequest( + sample, + time, + condition, + prediction_type="velocity", + model_kwargs={"endpoint_time": endpoint}, + ) + ) + + if self.config.jvp.method == "exact": + tangents = (dxt_dt.to(x_t.dtype), torch.ones_like(t), torch.zeros_like(r)) + rng_devices = [x_t.device] if x_t.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices), math_attention_for_forward_ad(x_t.device.type): + return torch.func.jvp(model_fn, (x_t, t, r), tangents)[1] + + work_t = t.double() + work_r = r.double() + epsilon = torch.full_like(work_t, self.config.jvp.epsilon) + forward_valid = work_t + epsilon <= 1.0 + backward_valid = (work_t - epsilon >= 0.0) & (work_t - epsilon > work_r) + central = forward_valid & backward_valid + forward = forward_valid & ~backward_valid + backward = ~forward_valid & backward_valid + + plus_t = work_t.clone() + minus_t = work_t.clone() + factor = torch.zeros_like(work_t) + plus_t[central] += epsilon[central] + minus_t[central] -= epsilon[central] + factor[central] = 0.5 / epsilon[central] + plus_t[forward] += epsilon[forward] + factor[forward] = 1.0 / epsilon[forward] + minus_t[backward] -= epsilon[backward] + factor[backward] = 1.0 / epsilon[backward] + + plus_x = x_t.double() + expand_time(plus_t - work_t, x_t.ndim) * dxt_dt.double() + minus_x = x_t.double() + expand_time(minus_t - work_t, x_t.ndim) * dxt_dt.double() + rng_devices = [x_t.device] if x_t.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices): + plus = model_fn(plus_x.to(x_t.dtype), plus_t.to(t.dtype), work_r.to(r.dtype)) + with torch.random.fork_rng(devices=rng_devices): + minus = model_fn(minus_x.to(x_t.dtype), minus_t.to(t.dtype), work_r.to(r.dtype)) + return (plus.double() - minus.double()) * expand_time(factor, plus.ndim) + + def _loss( + self, + velocity: Tensor, + velocity_jvp: Tensor, + dxt_dt: Tensor, + t: Tensor, + r: Tensor, + iteration: int, + ): + velocity = velocity.double() + velocity_jvp = velocity_jvp.double() + dxt_dt = dxt_dt.double() + delta = expand_time((t.double() - r.double()).clamp_min(0.0), velocity.ndim) + if self.config.tangent_warmup_steps: + warmup = min(1.0, iteration / self.config.tangent_warmup_steps) + else: + warmup = 1.0 + + if self.config.loss_type == "l2": + tangent = dxt_dt - warmup * delta * velocity_jvp + squared_error = (velocity - tangent).square().flatten(1).sum(dim=1) + weight = self._weight(squared_error) + loss = squared_error * weight + else: + tangent = dxt_dt - velocity.detach() - warmup * delta * velocity_jvp + if self.config.spatially_normalized_tangent: + tangent = tangent / (tangent[0].numel() ** 0.5) + tangent_norm = torch.linalg.vector_norm(tangent.flatten(1), dim=1) + weight = self._weight(tangent_norm) + target = (velocity + tangent * expand_time(weight, tangent.ndim)).detach() + loss = (velocity - target).square().flatten(1).sum(dim=1) + return loss, tangent, weight, warmup + + def _weight(self, value: Tensor) -> Tensor: + method, *arguments = self.config.norm_method.split("_") + if method == "poly" and len(arguments) == 1: + return (value + self.config.norm_constant).pow(-float(arguments[0])) + if method == "exp" and len(arguments) == 2: + constant, scale = map(float, arguments) + return constant * torch.exp(scale * value + self.config.norm_constant) + raise ValueError("MeanFlow loss.norm_method must be 'poly_' or 'exp__'.") diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/objective_factory.py b/lightx2v_train/lightx2v_train/trainers/consistency/objective_factory.py new file mode 100644 index 000000000..f68322753 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/objective_factory.py @@ -0,0 +1,29 @@ +import importlib + +from lightx2v_train.utils.registry import Register + +CONSISTENCY_OBJECTIVE_REGISTER = Register() + +_OBJECTIVE_MODULES = { + "cm": "lightx2v_train.trainers.consistency.cm", + "mean_flow": "lightx2v_train.trainers.consistency.mean_flow", + "meanflow": "lightx2v_train.trainers.consistency.mean_flow", + "pcm": "lightx2v_train.trainers.consistency.pcm", + "scm": "lightx2v_train.trainers.consistency.scm", + "tcm": "lightx2v_train.trainers.consistency.tcm", +} + + +def build_consistency_objective(config, path): + consistency_config = config["training"].get("consistency", {}) + if not isinstance(consistency_config, dict): + raise ValueError("training.consistency must be a mapping.") + algorithm = str(consistency_config.get("algorithm", "cm")).lower() + if algorithm not in CONSISTENCY_OBJECTIVE_REGISTER: + module_name = _OBJECTIVE_MODULES.get(algorithm) + if module_name is not None: + importlib.import_module(module_name) + if algorithm not in CONSISTENCY_OBJECTIVE_REGISTER: + available = ", ".join(sorted(CONSISTENCY_OBJECTIVE_REGISTER.keys())) + raise ValueError(f"Unknown consistency algorithm {algorithm!r}. Available algorithms: {available}") + return CONSISTENCY_OBJECTIVE_REGISTER[algorithm](config, path) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py new file mode 100644 index 000000000..e117a75ba --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +from torch import Tensor + +from .base import ( + ConsistencyBatch, + ConsistencyObjective, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, +) +from .cm import classifier_free_guidance +from .config import PCMConfig, PCMLossConfig, PCMSolverConfig +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER + + +@dataclass(frozen=True) +class PCMTimeState: + solver_index: Tensor + phase_index: Tensor + phase_start_index: Tensor + t: Tensor + r: Tensor + s: Tensor + + +class PCMTimeGrid: + """Discrete teacher-solver grid and equal-width PCM phase partition. + + Grid indices run from low to high noise, matching the reference SD3 + ``EulerSolver``. For an index ``i``, ``r`` is the adjacent lower-noise + solver time and ``s`` is the lower boundary shared by the whole phase. + """ + + def __init__(self, config: PCMSolverConfig): + self.config = config + + def sample(self, batch_size: int, scheduler, *, latent_hw, device) -> PCMTimeState: + current, previous = self.build_grid(scheduler, latent_hw=latent_hw, device=device) + solver_index = torch.randint(0, self.config.num_solver_steps, (batch_size,), device=device) + phase_starts = self.phase_start_indices(device) + phase_index = torch.bucketize(solver_index, phase_starts, right=True) - 1 + phase_start_index = phase_starts[phase_index] + return PCMTimeState( + solver_index=solver_index, + phase_index=phase_index, + phase_start_index=phase_start_index, + t=current[solver_index], + r=previous[solver_index], + s=previous[phase_start_index], + ) + + def build_grid(self, scheduler, *, latent_hw, device) -> tuple[Tensor, Tensor]: + steps = self.config.num_solver_steps + max_time = float(scheduler.max_t) + current = torch.arange(1, steps + 1, device=device, dtype=torch.float32) * (max_time / steps) + raw_boundary = scheduler.min_t if self.config.boundary_time is None else self.config.boundary_time + boundary = torch.tensor([raw_boundary], device=device, dtype=torch.float32) + + if scheduler.do_time_shift: + current = scheduler.time_shift(current, latent_hw=latent_hw) + if raw_boundary > 0.0: + boundary = scheduler.time_shift(boundary, latent_hw=latent_hw) + if not boundary.item() < current[0].item(): + raise ValueError( + "PCM boundary time must be smaller than the first solver time; " + f"got boundary={boundary.item():.6f}, first={current[0].item():.6f}." + ) + previous = torch.cat([boundary, current[:-1]]) + return current, previous + + def phase_start_indices(self, device) -> Tensor: + phases = torch.arange(self.config.num_phases, device=device, dtype=torch.int64) + return torch.div( + phases * self.config.num_solver_steps, + self.config.num_phases, + rounding_mode="floor", + ) + + +class PCMLoss: + """Elementwise pseudo-Huber/L2 loss used by the released PCM scripts.""" + + def __init__(self, config: PCMLossConfig): + self.config = config + self.dtype = torch.float64 if config.computation_dtype == "float64" else torch.float32 + + def __call__(self, prediction: Tensor, target: Tensor) -> Tensor: + difference = prediction.to(self.dtype) - target.to(self.dtype) + if self.config.distance == "l2": + elementwise = difference.square() + else: + constant = self.config.huber_constant + elementwise = torch.sqrt(difference.square() + constant**2) - constant + return elementwise.flatten(1).mean(dim=1) + + +def pcm_inference_sigmas(num_inference_steps: int, num_solver_steps: int) -> list[float]: + """Return the unshifted descending phase schedule used by PCM inference.""" + if not 1 <= num_inference_steps <= num_solver_steps: + raise ValueError("PCM inference steps must be in [1, num_solver_steps].") + indices = torch.div( + torch.arange(num_inference_steps, dtype=torch.int64) * num_solver_steps, + num_inference_steps, + rounding_mode="floor", + ) + return ((num_solver_steps - indices).float() / num_solver_steps).tolist() + + +@CONSISTENCY_OBJECTIVE_REGISTER("pcm") +class PCMObjective(ConsistencyObjective): + """Phased Consistency Model distillation for rectified-flow backbones.""" + + algorithm_name = "pcm" + requires_teacher = True + + def __init__(self, config, path: RectifiedFlowPath): + self.config = PCMConfig.from_mapping(config) + self.path = path + self.time_grid = PCMTimeGrid(self.config.solver) + self.loss_fn = PCMLoss(self.config.loss) + guidance_scale = self.config.teacher.guidance_scale + self.requires_negative_condition = guidance_scale is not None and guidance_scale != 1.0 + self.negative_prompt = self.config.teacher.negative_prompt + + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + del context + latent_hw = (clean.shape[-2], clean.shape[-1]) + time = self.time_grid.sample( + clean.shape[0], + scheduler, + latent_hw=latent_hw, + device=clean.device, + ) + return { + "noise": torch.randn_like(clean), + "solver_index": time.solver_index, + "phase_index": time.phase_index, + "phase_start_index": time.phase_start_index, + "t": time.t, + "r": time.r, + "s": time.s, + } + + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + del references + if teacher is None: + raise RuntimeError("PCM requires a frozen diffusion teacher denoiser.") + + clean = batch.clean + t = training_state["t"] + r = training_state["r"] + s = training_state["s"] + noisy_t = self.path.interpolate(clean, training_state["noise"], t) + + # Match dropout masks between the trainable prediction and the + # stop-gradient student target while leaving global RNG unchanged. + rng_devices = [clean.device] if clean.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices): + prediction_velocity = student.predict( + DenoiserRequest(noisy_t, t, batch.condition, prediction_type="velocity") + ) + prediction = self.path.euler_step(noisy_t, prediction_velocity, t, s) + + with torch.no_grad(): + teacher_velocity = self._teacher_velocity( + noisy_t, + t, + batch.condition, + batch.negative_condition, + teacher, + ) + noisy_r = self.path.euler_step(noisy_t, teacher_velocity, t, r) + target_velocity = student.predict( + DenoiserRequest(noisy_r, r, batch.condition, prediction_type="velocity") + ) + target = self.path.euler_step(noisy_r, target_velocity, r, s) + + per_sample_loss = self.loss_fn(prediction, target) + return ObjectiveOutput( + loss=per_sample_loss.mean(), + metrics={ + "pcm_unweighted": per_sample_loss.detach().mean(), + "pcm_t": t.detach().float().mean(), + "pcm_r": r.detach().float().mean(), + "pcm_s": s.detach().float().mean(), + "pcm_teacher_step": (t.float() - r.float()).detach().mean(), + "pcm_phase_span": (t.float() - s.float()).detach().mean(), + "pcm_phase": training_state["phase_index"].detach().float().mean(), + "pcm_solver_index": training_state["solver_index"].detach().float().mean(), + }, + ) + + @torch.no_grad() + def _teacher_velocity( + self, + sample: Tensor, + time: Tensor, + condition, + negative_condition, + teacher: ModelDenoiser, + ) -> Tensor: + conditional = teacher.predict( + DenoiserRequest(sample, time, condition, prediction_type="velocity") + ) + scale = self.config.teacher.guidance_scale + if scale is None or scale == 1.0: + return conditional + if negative_condition is None: + raise RuntimeError("PCM teacher CFG requires a negative condition.") + unconditional = teacher.predict( + DenoiserRequest(sample, time, negative_condition, prediction_type="velocity") + ) + return classifier_free_guidance( + conditional, + unconditional, + scale, + self.config.teacher.cfg_norm, + ) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/scm.py b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py new file mode 100644 index 000000000..ade52e30f --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +from typing import Mapping, Optional + +import torch +from torch import Tensor + +from .base import ( + ConsistencyBatch, + ConsistencyObjective, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, + expand_time, +) +from .cm import classifier_free_guidance +from .config import SCMConfig +from .jvp import math_attention_for_forward_ad +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER + + +@CONSISTENCY_OBJECTIVE_REGISTER("scm") +class SCMObjective(ConsistencyObjective): + """Continuous-time simplified Consistency Model (sCM) objective. + + The public model still predicts Qwen's rectified-flow velocity. This + objective performs FastGen's SNR-matched TrigFlow preconditioning at the + model boundary, keeping the backbone and its checkpoints interoperable. + """ + + algorithm_name = "scm" + model_capabilities = frozenset({"log_variance"}) + + def __init__(self, config, path: RectifiedFlowPath): + self.config = SCMConfig.from_mapping(config) + self.path = path + self.requires_teacher = self.config.mode == "cd" + guidance_scale = self.config.teacher.guidance_scale + self.requires_negative_condition = self.requires_teacher and guidance_scale is not None and guidance_scale != 1.0 + self.negative_prompt = self.config.teacher.negative_prompt + + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + latent_hw = (clean.shape[-2], clean.shape[-1]) + t = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() + return { + "noise": torch.randn_like(clean) * self.config.sigma_data, + "t": t, + "iteration": torch.tensor(context.iteration, device=clean.device, dtype=torch.int64), + } + + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + del references + clean = batch.clean + z = training_state["noise"] + t = training_state["t"].float() + sigma_data = self.config.sigma_data + + alpha = 1.0 - t + sigma = t + t_hat = torch.atan2(sigma, alpha * sigma_data) + alpha_hat = torch.cos(t_hat) + sigma_hat = torch.sin(t_hat) + x_hat = expand_time(alpha_hat, clean.ndim).to(clean) * clean + expand_time(sigma_hat, clean.ndim).to(clean) * z + + if self.config.mode == "ct": + dxt_dt = -expand_time(sigma_hat, clean.ndim).to(clean) * clean + expand_time(alpha_hat, clean.ndim).to(clean) * z + else: + if teacher is None: + raise RuntimeError("sCM consistency distillation requires a frozen teacher denoiser.") + with torch.no_grad(): + teacher_flow, _, _ = self._predict_trig_flow(teacher, x_hat, t_hat, batch.condition) + dxt_dt = sigma_data * teacher_flow + guidance_scale = self.config.teacher.guidance_scale + if guidance_scale is not None and guidance_scale != 1.0: + if batch.negative_condition is None: + raise RuntimeError("sCM teacher CFG requires a negative condition.") + negative_flow, _, _ = self._predict_trig_flow( + teacher, + x_hat, + t_hat, + batch.negative_condition, + ) + guided_flow = classifier_free_guidance( + teacher_flow, + negative_flow, + guidance_scale, + self.config.teacher.cfg_norm, + ) + dxt_dt = sigma_data * guided_flow + + flow, original_t, x0_prediction = self._predict_trig_flow(student, x_hat, t_hat, batch.condition) + log_variance = student.predict_log_variance(original_t).reshape(-1) + flow_jvp = self._jvp(clean, z, x_hat, t_hat, dxt_dt, batch.condition, student) + + loss, unweighted, tangent, warmup = self._loss( + flow, + flow_jvp, + x_hat, + dxt_dt, + log_variance, + sigma, + t_hat, + int(training_state["iteration"].item()), + ) + return ObjectiveOutput( + loss=loss.mean(), + metrics={ + "scm_unweighted": unweighted.detach().mean(), + "scm_log_variance": log_variance.detach().mean(), + "scm_t": t.detach().mean(), + "scm_t_hat": t_hat.detach().mean(), + "scm_jvp_norm": torch.linalg.vector_norm(flow_jvp.flatten(1), dim=1).detach().mean(), + "scm_tangent_norm": torch.linalg.vector_norm(tangent.flatten(1), dim=1).detach().mean(), + "scm_warmup": warmup, + "scm_x0_norm": torch.linalg.vector_norm(x0_prediction.detach().flatten(1), dim=1).mean(), + }, + ) + + def _predict_trig_flow(self, denoiser: ModelDenoiser, x_hat: Tensor, t_hat: Tensor, condition): + sigma_data = self.config.sigma_data + tangent = sigma_data * torch.tan(t_hat.double()) + original_t = (tangent / (1.0 + tangent)).to(t_hat.dtype) + coefficient = torch.sqrt((1.0 - original_t.double()).square() + (original_t.double() / sigma_data).square()).to(x_hat.dtype) + original_sample = x_hat * expand_time(coefficient, x_hat.ndim) + x0_prediction = denoiser.predict( + DenoiserRequest( + sample=original_sample, + time=original_t, + condition=condition, + prediction_type="x0", + ) + ) + denominator = expand_time(torch.sin(t_hat), x_hat.ndim).to(x_hat).clamp_min(1e-6) + trig_velocity = (expand_time(torch.cos(t_hat), x_hat.ndim).to(x_hat) * x_hat - x0_prediction) / denominator + return trig_velocity / sigma_data, original_t, x0_prediction + + @torch.no_grad() + def _jvp( + self, + clean: Tensor, + z: Tensor, + x_hat: Tensor, + t_hat: Tensor, + dxt_dt: Tensor, + condition, + student: ModelDenoiser, + ) -> Tensor: + def model_fn(sample, time): + return self._predict_trig_flow(student, sample, time, condition)[0] + + if self.config.jvp.method == "exact": + v_t = torch.cos(t_hat) * torch.sin(t_hat) + v_x = expand_time(v_t, dxt_dt.ndim).to(dxt_dt) * dxt_dt + rng_devices = [x_hat.device] if x_hat.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices), math_attention_for_forward_ad(x_hat.device.type): + return torch.func.jvp(model_fn, (x_hat, t_hat), (v_x, v_t))[1] + + work_t = t_hat.double().clamp(1e-5, torch.pi / 2 - 1e-5) + epsilon = (self.config.jvp.epsilon * work_t.abs()).clamp_min(1e-6) + plus_t = (work_t + epsilon).clamp_max(torch.pi / 2 - 1e-5) + minus_t = (work_t - epsilon).clamp_min(1e-5) + plus_x = expand_time(torch.cos(plus_t), clean.ndim) * clean.double() + expand_time(torch.sin(plus_t), clean.ndim) * z.double() + minus_x = expand_time(torch.cos(minus_t), clean.ndim) * clean.double() + expand_time(torch.sin(minus_t), clean.ndim) * z.double() + rng_devices = [x_hat.device] if x_hat.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices): + plus = model_fn(plus_x.to(x_hat.dtype), plus_t.to(t_hat.dtype)) + with torch.random.fork_rng(devices=rng_devices): + minus = model_fn(minus_x.to(x_hat.dtype), minus_t.to(t_hat.dtype)) + v_t = torch.cos(work_t) * torch.sin(work_t) + factor = expand_time(v_t / (plus_t - minus_t), plus.ndim) + return (plus.double() - minus.double()) * factor + + def _loss( + self, + flow: Tensor, + flow_jvp: Tensor, + x_hat: Tensor, + dxt_dt: Tensor, + log_variance: Tensor, + sigma: Tensor, + t_hat: Tensor, + iteration: int, + ): + flow = flow.double() + detached_flow = flow.detach() + flow_jvp = flow_jvp.double() + x_hat = x_hat.double() + dxt_dt = dxt_dt.double() + alpha_hat = torch.cos(t_hat.double()) + sigma_hat = torch.sin(t_hat.double()) + if self.config.tangent_warmup_steps: + warmup = min(1.0, iteration / self.config.tangent_warmup_steps) + else: + warmup = 1.0 + + g1 = -expand_time(alpha_hat.square(), flow.ndim) * (self.config.sigma_data * detached_flow - dxt_dt) + g2 = -(expand_time(alpha_hat * sigma_hat, flow.ndim) * x_hat + self.config.sigma_data * flow_jvp) + tangent = g1 + warmup * g2 + tangent_norm = torch.linalg.vector_norm(tangent.flatten(1), dim=1) + if self.config.spatially_normalized_tangent: + tangent_norm = tangent_norm / (tangent[0].numel() ** 0.5) + tangent = tangent / expand_time(tangent_norm + self.config.tangent_warmup_constant, tangent.ndim) + + unweighted = (flow - detached_flow - tangent).square().flatten(1).mean(dim=1) + prior_weight = sigma.double().clamp_min(self.config.min_denominator).reciprocal() if self.config.prior_weighting else torch.ones_like(sigma, dtype=torch.float64) + dimension = x_hat[0].numel() if self.config.normalize_by_numel else 1.0 + log_variance = log_variance.double() + loss = prior_weight * torch.exp(-log_variance) * unweighted / dimension + log_variance + return loss, unweighted, tangent, warmup diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py new file mode 100644 index 000000000..21b984523 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from typing import Mapping, Optional + +import torch +import torch.distributed as dist +from torch import Tensor + +from .base import ( + ConsistencyBatch, + ConsistencyStepContext, + DenoiserRequest, + ModelDenoiser, + ObjectiveOutput, + RectifiedFlowPath, + ReferenceModelSpec, + expand_time, +) +from .cm import CMObjective +from .config import TCMConfig +from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER + + +def _distributed_mask_summary(mask: Tensor) -> tuple[bool, bool]: + """Return global ``(all, any)`` without creating FSDP branch skew.""" + all_value = torch.tensor(bool(mask.all()), device=mask.device, dtype=torch.int32) + any_value = torch.tensor(bool(mask.any()), device=mask.device, dtype=torch.int32) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(all_value, op=dist.ReduceOp.MIN) + dist.all_reduce(any_value, op=dist.ReduceOp.MAX) + return bool(all_value.item()), bool(any_value.item()) + + +class TwoStageDenoiser: + """Route low times through frozen stage 1 and high times through stage 2.""" + + def __init__(self, stage1: ModelDenoiser, stage2: ModelDenoiser, transition_time: float): + self.stage1 = stage1 + self.stage2 = stage2 + self.transition_time = transition_time + + def predict(self, request: DenoiserRequest) -> Tensor: + second_stage = request.time >= self.transition_time + all_second, any_second = _distributed_mask_summary(second_stage) + if all_second: + return self.stage2.predict(request) + + rng_devices = [request.sample.device] if request.sample.device.type == "cuda" else [] + with torch.random.fork_rng(devices=rng_devices), torch.no_grad(): + stage1_prediction = self.stage1.predict(request) + if not any_second: + return stage1_prediction + + stage2_prediction = self.stage2.predict(request) + mask = expand_time(second_stage, stage2_prediction.ndim) + return torch.where(mask, stage2_prediction, stage1_prediction) + + +@CONSISTENCY_OBJECTIVE_REGISTER("tcm") +class TCMObjective(CMObjective): + """Stage-2 Consistency Model training with a frozen stage-1 boundary.""" + + algorithm_name = "tcm" + + def __init__(self, config, path: RectifiedFlowPath): + super().__init__(config, path) + self.tcm_config = TCMConfig.from_mapping(config) + + @property + def reference_model_specs(self) -> tuple[ReferenceModelSpec, ...]: + return ( + ReferenceModelSpec( + role="stage1", + checkpoint=self.tcm_config.stage1_checkpoint, + # FastGen intentionally keeps stage 1 in train mode so its + # dropout RNG can be matched to the stage-2 forward. + training_mode=True, + ), + ) + + @property + def student_initialization_checkpoint(self) -> str: + return self.tcm_config.stage1_checkpoint + + def sample_training_state( + self, + clean: Tensor, + scheduler, + context: ConsistencyStepContext, + ) -> Mapping[str, Tensor]: + batch_size = clean.shape[0] + latent_hw = (clean.shape[-2], clean.shape[-1]) + t = scheduler.sample_timestep_or_sigma(batch_size, latent_hw=latent_hw).to(clean.device).float() + t = t.clamp_min(self.tcm_config.transition_time + self.config.time_pair.safety_epsilon) + + boundary_count = int(batch_size * self.tcm_config.boundary_probability) + if boundary_count: + permutation = torch.randperm(batch_size, device=clean.device) + boundary_mask = torch.zeros(batch_size, device=clean.device, dtype=torch.bool) + boundary_mask[permutation[:boundary_count]] = True + else: + boundary_mask = torch.rand(batch_size, device=clean.device) < self.tcm_config.boundary_probability + t = torch.where( + boundary_mask, + torch.full_like(t, self.tcm_config.transition_time + self.config.time_pair.safety_epsilon), + t, + ) + + pair = self.time_pair_sampler.sample(t, context) + return { + "noise": torch.randn_like(clean), + "t": pair.t, + "r": pair.r, + "boundary_mask": boundary_mask, + "ratio": torch.tensor(pair.ratio, device=clean.device, dtype=torch.float32), + "stage": torch.tensor(pair.stage, device=clean.device, dtype=torch.float32), + } + + def compute( + self, + batch: ConsistencyBatch, + training_state: Mapping[str, Tensor], + student: ModelDenoiser, + teacher: Optional[ModelDenoiser] = None, + references: Optional[Mapping[str, ModelDenoiser]] = None, + ) -> ObjectiveOutput: + if references is None or "stage1" not in references: + raise RuntimeError("TCM requires the frozen 'stage1' consistency model.") + + two_stage = TwoStageDenoiser( + references["stage1"], + student, + self.tcm_config.transition_time, + ) + terms = self.compute_loss_terms(batch, training_state, two_stage, teacher) + boundary_mask = training_state["boundary_mask"].bool() + zero = terms.weighted.sum() * 0.0 + regular_loss = terms.weighted[~boundary_mask].mean() if (~boundary_mask).any() else zero + boundary_loss = terms.weighted[boundary_mask].mean() if boundary_mask.any() else zero + regular_unweighted = terms.unweighted[~boundary_mask].mean() if (~boundary_mask).any() else zero.detach() + loss = regular_loss + self.tcm_config.boundary_weight * boundary_loss + + return ObjectiveOutput( + loss=loss, + metrics={ + "tcm_regular": regular_loss.detach(), + "tcm_boundary": boundary_loss.detach(), + "tcm_unweighted": regular_unweighted.detach(), + "tcm_boundary_fraction": boundary_mask.float().mean(), + "tcm_t": training_state["t"].detach().mean(), + "tcm_r": training_state["r"].detach().mean(), + "tcm_ratio": training_state["ratio"].detach(), + "tcm_stage": training_state["stage"].detach(), + }, + ) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py new file mode 100644 index 000000000..d3c1bd4f9 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import copy +import os + +import torch +from loguru import logger + +from lightx2v_train.model_zoo import build_model +from lightx2v_train.runtime.distributed import get_data_parallel_world_size +from lightx2v_train.runtime.parallel import apply_parallel +from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value +from lightx2v_train.utils.registry import TRAINER_REGISTER + +from ..flow import FlowMatchingTrainer +from .base import ( + ConsistencyBatch, + ConsistencyStepContext, + ModelDenoiser, + RectifiedFlowPath, +) +from .objective_factory import build_consistency_objective + + +@TRAINER_REGISTER("cm") +@TRAINER_REGISTER("consistency") +class ConsistencyTrainer(FlowMatchingTrainer): + """Shared trainer shell for consistency-model objectives. + + The trainer owns data encoding, distributed synchronization, optimizer and + checkpoint state, and the optional frozen teacher. Algorithm-specific + mathematics lives behind ``ConsistencyObjective``. + """ + + trainer_name = "consistency" + + def __init__(self, config): + super().__init__(config) + self.path = RectifiedFlowPath() + self.objective = build_consistency_objective(config, self.path) + self.teacher_model = None + self.teacher_denoiser = None + self.reference_models = {} + self.reference_denoisers = {} + self._setup_resume_checkpoint = None + + def setup(self, resume_ckpt_path=None): + self._setup_resume_checkpoint = resume_ckpt_path + super().setup(resume_ckpt_path=resume_ckpt_path) + self.student_denoiser = ModelDenoiser(self.model, self.path) + + if self.objective.requires_teacher: + self.teacher_model = self._build_frozen_teacher() + self.teacher_denoiser = ModelDenoiser(self.teacher_model, self.path) + + for spec in self.objective.reference_model_specs: + reference_model = self._build_frozen_reference(spec) + self.reference_models[spec.role] = reference_model + self.reference_denoisers[spec.role] = ModelDenoiser(reference_model, self.path) + + logger.info( + "[train] consistency algorithm={} mode={} teacher={}", + self.objective.algorithm_name, + getattr(getattr(self.objective, "config", None), "mode", "custom"), + self.objective.requires_teacher, + ) + + def _setup_trainable_model(self, model): + model.configure_consistency_model(self.objective.model_capabilities) + super()._setup_trainable_model(model) + model.set_consistency_modules_trainable() + + initialization = self.objective.student_initialization_checkpoint + if initialization is not None and self._setup_resume_checkpoint is None: + self._load_initial_model_weights(model, initialization) + # Loading weights does not change requires_grad, but keeping this + # call here makes that lifecycle guarantee explicit for new models. + model.set_consistency_modules_trainable() + + def _restore_trainable_model(self, model): + super()._restore_trainable_model(model) + model.set_consistency_modules_trainable() + + def _load_initial_model_weights(self, model, checkpoint, *, role="student"): + if not os.path.isdir(checkpoint): + raise RuntimeError(f"Consistency initialization checkpoint does not exist: {checkpoint}") + self._load_model_weights(model, checkpoint) + logger.info("[train] initialized consistency {} from {}", role, checkpoint) + + def _build_frozen_teacher(self): + role_keys = { + "fake", + "fake_2", + "fake_low_high", + "fake_real", + "fake_real_high", + "fake_real_low", + "student", + "student_2", + "stage1", + "teacher", + "teacher_2", + } + base_model_config = {key: copy.deepcopy(value) for key, value in self.model_config.items() if key not in role_keys} + teacher_override = self.model_config.get("teacher", {}) + if not isinstance(teacher_override, dict): + raise ValueError("model.teacher must be a mapping when provided.") + + teacher_config = copy.deepcopy(self.config) + teacher_config["model"] = base_model_config + teacher_config["model"].update(copy.deepcopy(teacher_override)) + teacher_model = build_model(teacher_config) + teacher_model.load_components(transformer_only=True, reference_model=self.model) + teacher_model.denoiser_module().requires_grad_(False) + teacher_model.set_denoiser_eval() + apply_parallel(teacher_model, self.config) + teacher_model.set_denoiser_eval() + logger.info( + "[train] consistency teacher model={} path={}", + teacher_config["model"]["name"], + teacher_config["model"]["pretrained_model_name_or_path"], + ) + return teacher_model + + def _build_frozen_reference(self, spec): + role_keys = { + "fake", + "fake_2", + "fake_low_high", + "fake_real", + "fake_real_high", + "fake_real_low", + "student", + "student_2", + "stage1", + "teacher", + "teacher_2", + } + model_config = {key: copy.deepcopy(value) for key, value in self.model_config.items() if key not in role_keys} + override = self.model_config.get(spec.role, {}) + if not isinstance(override, dict): + raise ValueError(f"model.{spec.role} must be a mapping when provided.") + model_config.update(copy.deepcopy(override)) + + reference_config = copy.deepcopy(self.config) + reference_config["model"] = model_config + reference_model = build_model(reference_config) + reference_model.load_components(transformer_only=True, reference_model=self.model) + if self.train_type == "lora": + reference_model.add_lora(self.lora_rank, self.lora_alpha, self.lora_target_modules) + self._load_initial_model_weights(reference_model, spec.checkpoint, role=spec.role) + reference_model.denoiser_module().requires_grad_(False) + if spec.training_mode: + reference_model.denoiser_module().train() + else: + reference_model.set_denoiser_eval() + apply_parallel(reference_model, self.config) + if spec.training_mode: + reference_model.denoiser_module().train() + else: + reference_model.set_denoiser_eval() + logger.info( + "[train] consistency reference role={} model={} checkpoint={}", + spec.role, + reference_config["model"]["name"], + spec.checkpoint, + ) + return reference_model + + def compute_loss_on_sample(self, sample): + with torch.no_grad(): + clean = self.model.encode_to_latent(sample) + clean = broadcast_sequence_parallel_value(clean) + condition = self.model.encode_condition(sample) + condition = broadcast_sequence_parallel_value(condition) + + negative_condition = None + if self.objective.requires_negative_condition: + negative_condition = self._encode_negative_condition(sample, clean.shape[0]) + negative_condition = broadcast_sequence_parallel_value(negative_condition) + + context = ConsistencyStepContext( + iteration=int(getattr(self, "current_train_iteration", 0)), + global_batch_size=(clean.shape[0] * get_data_parallel_world_size() * self.gradient_accumulation_iters), + ) + training_state = self.objective.sample_training_state( + clean, + self.noise_scheduler, + context, + ) + training_state = broadcast_sequence_parallel_value(training_state) + + output = self.objective.compute( + ConsistencyBatch( + clean=clean, + condition=condition, + negative_condition=negative_condition, + ), + training_state, + self.student_denoiser, + self.teacher_denoiser, + self.reference_denoisers, + ) + return {"loss": output.loss, "metrics": output.metrics} + + def _encode_negative_condition(self, sample, batch_size): + conditioning = sample.get("conditioning", {}) + prompt = conditioning.get("prompt", "") + negative_prompt = conditioning.get("negative_prompt") + + if negative_prompt is None: + values = [self.objective.negative_prompt] * batch_size + elif isinstance(negative_prompt, str): + values = [negative_prompt] * batch_size + else: + values = list(negative_prompt) + if len(values) == 1 and batch_size > 1: + values *= batch_size + elif len(values) != batch_size: + raise ValueError(f"Expected {batch_size} negative prompts, got {len(values)}.") + + fallback = self.objective.negative_prompt or " " + values = [value if isinstance(value, str) and value.strip() else fallback for value in values] + encoded_prompt = values[0] if isinstance(prompt, str) and batch_size == 1 else values + + negative_sample = dict(sample) + negative_sample["conditioning"] = dict(conditioning) + negative_sample["conditioning"]["prompt"] = encoded_prompt + return self.model.encode_condition(negative_sample) diff --git a/lightx2v_train/lightx2v_train/trainers/flow.py b/lightx2v_train/lightx2v_train/trainers/flow.py index bd05a1a60..2fddec934 100644 --- a/lightx2v_train/lightx2v_train/trainers/flow.py +++ b/lightx2v_train/lightx2v_train/trainers/flow.py @@ -48,6 +48,9 @@ def compute_loss_on_sample(self, sample): def train(self): resume_ckpt_path, current_iter = self._resolve_resume() + # Objectives with a curriculum (for example CM's shrinking time + # interval) need the optimizer-step index, including after resume. + self.current_train_iteration = current_iter self.setup(resume_ckpt_path=resume_ckpt_path) if is_main_process(): os.makedirs(self.output_train_dir, exist_ok=True) @@ -111,6 +114,7 @@ def train(self): self.optimizer.zero_grad() current_iter += 1 + self.current_train_iteration = current_iter display_loss = reduce_mean(running_loss) current_lr = self.lr_scheduler.get_last_lr()[0] if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index cfaa4ccf6..d202c454f 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -90,6 +90,8 @@ def merge(self, other_register): } _TRAINER_MODULES = { + "cm": "lightx2v_train.trainers.consistency.trainer", + "consistency": "lightx2v_train.trainers.consistency.trainer", "dmd": "lightx2v_train.trainers.dmd.trainer", "dopsd": "lightx2v_train.trainers.dopsd", "fastwam": "lightx2v_train.trainers.fastwam", diff --git a/lightx2v_train/scripts/run_qwen_image_t2i_lora.sh b/lightx2v_train/scripts/run_qwen_image_t2i_lora.sh index 2eca39f4a..11703b4c0 100755 --- a/lightx2v_train/scripts/run_qwen_image_t2i_lora.sh +++ b/lightx2v_train/scripts/run_qwen_image_t2i_lora.sh @@ -1,6 +1,7 @@ #!/bin/bash export CUDA_VISIBLE_DEVICES=7 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun \ --standalone \ From 53ac25e5a05759a234e27bdd4f7cf40afa72dc06 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 3 Aug 2026 18:09:54 +0000 Subject: [PATCH 02/34] update --- .../trainers/consistency/config.py | 5 +---- .../trainers/consistency/pcm.py | 21 +++++-------------- 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/config.py b/lightx2v_train/lightx2v_train/trainers/consistency/config.py index 5b6474b92..df35d1788 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/config.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/config.py @@ -290,10 +290,7 @@ def from_mapping(cls, config): consistency = _mapping(config["training"].get("consistency"), "training.consistency") mode = str(consistency.get("mode", "cd")).lower() if mode != "cd": - raise ValueError( - "The published PCM algorithm is distillation-only; " - "set training.consistency.mode to 'cd'." - ) + raise ValueError("The published PCM algorithm is distillation-only; set training.consistency.mode to 'cd'.") return cls( mode=mode, solver=PCMSolverConfig.from_mapping(consistency.get("solver")), diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py index e117a75ba..ea21cb6ed 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -68,10 +68,7 @@ def build_grid(self, scheduler, *, latent_hw, device) -> tuple[Tensor, Tensor]: if raw_boundary > 0.0: boundary = scheduler.time_shift(boundary, latent_hw=latent_hw) if not boundary.item() < current[0].item(): - raise ValueError( - "PCM boundary time must be smaller than the first solver time; " - f"got boundary={boundary.item():.6f}, first={current[0].item():.6f}." - ) + raise ValueError(f"PCM boundary time must be smaller than the first solver time; got boundary={boundary.item():.6f}, first={current[0].item():.6f}.") previous = torch.cat([boundary, current[:-1]]) return current, previous @@ -175,9 +172,7 @@ def compute( # stop-gradient student target while leaving global RNG unchanged. rng_devices = [clean.device] if clean.device.type == "cuda" else [] with torch.random.fork_rng(devices=rng_devices): - prediction_velocity = student.predict( - DenoiserRequest(noisy_t, t, batch.condition, prediction_type="velocity") - ) + prediction_velocity = student.predict(DenoiserRequest(noisy_t, t, batch.condition, prediction_type="velocity")) prediction = self.path.euler_step(noisy_t, prediction_velocity, t, s) with torch.no_grad(): @@ -189,9 +184,7 @@ def compute( teacher, ) noisy_r = self.path.euler_step(noisy_t, teacher_velocity, t, r) - target_velocity = student.predict( - DenoiserRequest(noisy_r, r, batch.condition, prediction_type="velocity") - ) + target_velocity = student.predict(DenoiserRequest(noisy_r, r, batch.condition, prediction_type="velocity")) target = self.path.euler_step(noisy_r, target_velocity, r, s) per_sample_loss = self.loss_fn(prediction, target) @@ -218,17 +211,13 @@ def _teacher_velocity( negative_condition, teacher: ModelDenoiser, ) -> Tensor: - conditional = teacher.predict( - DenoiserRequest(sample, time, condition, prediction_type="velocity") - ) + conditional = teacher.predict(DenoiserRequest(sample, time, condition, prediction_type="velocity")) scale = self.config.teacher.guidance_scale if scale is None or scale == 1.0: return conditional if negative_condition is None: raise RuntimeError("PCM teacher CFG requires a negative condition.") - unconditional = teacher.predict( - DenoiserRequest(sample, time, negative_condition, prediction_type="velocity") - ) + unconditional = teacher.predict(DenoiserRequest(sample, time, negative_condition, prediction_type="velocity")) return classifier_free_guidance( conditional, unconditional, From 7b3385a256e7c43e3ca641981c8326250087d97e Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 10 Aug 2026 07:55:24 +0000 Subject: [PATCH 03/34] big refactor --- .../infer/lingbot_video_t2v_dmd_lora.yaml | 1 - lightx2v_train/configs/infer/wan.yaml | 1 - .../infer/wan2_1_t2v_14b_tf_chunkwise_ar.yaml | 1 - .../wan2_1_t2v_1_3b_tf_chunkwise_ar.yaml | 1 - .../configs/infer/wan2_2_ti2v_5b_tf_ar.yaml | 1 - .../configs/train/dmd/flux2_dev_cdm_lora.yaml | 9 +- .../configs/train/dmd/flux2_dev_dmd_lora.yaml | 9 +- .../train/dmd/flux2_klein_dmd_lora.yaml | 13 +- .../train/dmd/lingbot_video_t2v_dmd_lora.yaml | 4 +- .../dmd/lingbot_video_t2v_droid_dmd_lora.yaml | 4 +- .../train/dmd/longcat_image_cdm_lora.yaml | 13 +- .../train/dmd/longcat_image_dmd_lora.yaml | 13 +- .../train/dmd/ltx_t2av_ar_dmd_lora.yaml | 3 +- .../configs/train/dmd/ltx_t2av_dmd_lora.yaml | 3 +- .../train/dmd/qwen_image_cdm_lora.yaml | 13 +- .../train/dmd/qwen_image_dmd_lora.yaml | 13 +- .../train/dmd/wan2_1_t2v_14b_ar_dmd.yaml | 4 +- .../train/dmd/wan2_1_t2v_14b_ar_dmd_lora.yaml | 4 +- .../configs/train/dmd/wan2_1_t2v_14b_dmd.yaml | 4 +- .../train/dmd/wan2_1_t2v_14b_dmd_lora.yaml | 4 +- .../train/dmd/wan2_1_t2v_1_3b_ar_dmd.yaml | 4 +- .../train/dmd/wan2_1_t2v_1_3b_dmd.yaml | 6 +- .../train/dmd/wan2_2_ti2v_5b_ar_dmd.yaml | 4 +- .../configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml | 4 +- .../train/dmd/wan2_2_ti2v_5b_dmd_multi.yaml | 4 +- .../libero_plus_uncond_2cam224_4gpu.yaml | 5 +- .../libero_plus_uncond_2cam224_8gpu.yaml | 5 +- .../train/fastwam/libero_uncond_2cam224.yaml | 5 +- .../configs/train/flow/flux2_dev_lora.yaml | 2 +- .../configs/train/flow/flux2_klein_lora.yaml | 2 +- .../configs/train/flow/flux2_klein_t2i.yaml | 2 +- .../train/flow/longcat_image_lora.yaml | 2 +- .../configs/train/flow/longcat_image_t2i.yaml | 2 +- .../configs/train/flow/ltx_t2av_full.yaml | 3 +- .../configs/train/flow/ltx_t2av_lora.yaml | 3 +- .../configs/train/flow/ltx_t2av_train.yaml | 3 +- .../train/flow/qwen_image_edit_2511_lora.yaml | 2 +- .../flow/qwen_image_edit_2511_lora_dist.yaml | 2 +- .../configs/train/flow/qwen_image_lora.yaml | 2 +- .../train/flow/qwen_image_lora_dist.yaml | 2 +- .../train/flow/qwen_image_lora_dp.yaml | 2 +- .../configs/train/flow/qwen_image_t2i.yaml | 2 +- .../train/flow/wan2_1_t2v_14b_full.yaml | 4 +- .../train/flow/wan2_1_t2v_14b_lora.yaml | 4 +- .../train/flow/wan2_1_t2v_1_3b_full.yaml | 4 +- .../train/flow/wan2_1_t2v_1_3b_lora.yaml | 4 +- .../wan2_1_t2v_1_3b_phased_dmd.yaml | 6 +- .../wan2_2_t2v_a14b_phased_dmd.yaml | 4 +- .../train/sgmd/wan2_1_t2v_14b_sgmd.yaml | 4 +- .../train/sgmd/wan2_1_t2v_1_3b_sgmd.yaml | 6 +- .../train/sgmd/wan2_2_ti2v_5b_sgmd.yaml | 4 +- .../tf/ltx_t2av_teacher_forcing_lmdb.yaml | 3 +- .../tf/wan2_1_t2v_14b_tf_chunkwise_full.yaml | 2 - .../tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml | 2 - .../tf/wan2_2_ti2v_5b_tf_cache_96f_sp4.yaml | 2 - .../wan2_2_ti2v_5b_tf_chunkwise_full_sp4.yaml | 2 - lightx2v_train/data_process/README.md | 75 +- .../lightx2v_train/data/libero/dataset.py | 2 +- .../lightx2v_train/data/libero/preparation.py | 6 +- lightx2v_train/lightx2v_train/data/utils.py | 6 + .../lightx2v_train/data/video_dataset.py | 200 ++- .../model_capabilities/__init__.py | 49 + .../lightx2v_train/model_capabilities/base.py | 94 ++ .../model_capabilities/training.py | 423 ++++++ .../lightx2v_train/model_zoo/__init__.py | 37 +- .../lightx2v_train/model_zoo/base.py | 33 +- .../model_zoo/capability_adapters/__init__.py | 1 + .../model_zoo/capability_adapters/common.py | 492 ++++++ .../model_zoo/flux2/__init__.py | 1 + .../flux2/capability_adapters/__init__.py | 1 + .../flux2_dopsd_capability.py | 201 +++ .../model_zoo/{ => flux2}/flux2_dev.py | 19 +- .../model_zoo/{ => flux2}/flux2_klein.py | 30 +- .../model_zoo/longcat_image/__init__.py | 1 + .../{ => longcat_image}/longcat_image.py | 19 +- .../lightx2v_train/model_zoo/ltx/__init__.py | 1 + .../ltx/capability_adapters/__init__.py | 1 + .../ltx/capability_adapters/common.py | 21 + ..._autoregressive_distillation_capability.py | 387 +++++ .../ltx_distillation_capability.py | 440 ++++++ .../ltx_flow_matching_capability.py | 204 +++ .../ltx_teacher_forcing_capability.py | 357 +++++ .../model_zoo/{ => ltx}/ltx_t2av.py | 36 +- .../native/wan/fastwam/action_dit.py | 4 +- .../model_zoo/native/wan/fastwam/model.py | 17 +- .../model_zoo/native/wan/fastwam/video_dit.py | 4 +- .../model_zoo/qwen_image/__init__.py | 1 + .../model_zoo/{ => qwen_image}/qwen_image.py | 23 +- .../{ => qwen_image}/qwen_image_edit.py | 11 +- .../lightx2v_train/model_zoo/wan/__init__.py | 1 + .../wan/capability_adapters/__init__.py | 1 + ..._autoregressive_distillation_capability.py | 303 ++++ .../wan_distillation_capability.py | 123 ++ .../wan_teacher_forcing_capability.py | 122 ++ .../wan_world_action_capability.py | 40 + .../model_zoo/{ => wan}/lingbot_video.py | 39 +- .../model_zoo/{ => wan}/wan_fastwam.py | 16 +- .../model_zoo/{ => wan}/wan_t2v.py | 58 +- .../model_zoo/{ => wan}/wan_ti2v_5b.py | 12 +- .../schedulers/dmd_scheduler.py | 16 +- .../schedulers/flow_matching.py | 36 +- .../lightx2v_train/trainers/__init__.py | 76 +- .../lightx2v_train/trainers/base.py | 97 +- .../trainers/consistency/__init__.py | 4 +- .../trainers/consistency/base.py | 38 +- .../lightx2v_train/trainers/consistency/cm.py | 19 +- .../trainers/consistency/mean_flow.py | 22 +- .../trainers/consistency/pcm.py | 17 +- .../trainers/consistency/scm.py | 18 +- .../trainers/consistency/tcm.py | 23 +- .../trainers/consistency/trainer.py | 119 +- .../lightx2v_train/trainers/dmd/__init__.py | 22 +- .../trainers/dmd/autoregressive_dmd.py | 72 + .../lightx2v_train/trainers/dmd/checkpoint.py | 62 +- .../lightx2v_train/trainers/dmd/config.py | 19 +- .../trainers/dmd/ltx_trainer.py | 675 --------- .../lightx2v_train/trainers/dmd/runtime.py | 459 ++++++ .../lightx2v_train/trainers/dmd/trainer.py | 1323 +++++++++-------- .../trainers/dmd/video_ar_trainer.py | 253 ---- .../trainers/dmd/video_trainer.py | 826 ---------- .../lightx2v_train/trainers/dopsd.py | 230 ++- .../lightx2v_train/trainers/flow.py | 293 ---- .../lightx2v_train/trainers/flow_matching.py | 136 ++ .../trainers/phased_dmd/__init__.py | 5 +- .../trainers/phased_dmd/checkpoint.py | 61 +- .../trainers/phased_dmd/config.py | 10 +- .../trainers/phased_dmd/math.py | 6 +- .../trainers/phased_dmd/roles.py | 101 +- .../trainers/phased_dmd/rollout.py | 87 +- .../trainers/phased_dmd/trainer.py | 23 +- .../lightx2v_train/trainers/sgmd.py | 27 +- .../trainers/teacher_forcing.py | 75 + lightx2v_train/lightx2v_train/trainers/tf.py | 426 ------ .../trainers/{fastwam.py => world_action.py} | 48 +- lightx2v_train/lightx2v_train/tricks/cdm.py | 11 - .../lightx2v_train/tricks/diversity.py | 4 - .../lightx2v_train/tricks/real_data_fake.py | 17 +- .../utils/image_size_buckets.py | 90 ++ .../lightx2v_train/utils/registry.py | 55 +- 139 files changed, 6086 insertions(+), 3964 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_capabilities/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_capabilities/base.py create mode 100644 lightx2v_train/lightx2v_train/model_capabilities/training.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py rename lightx2v_train/lightx2v_train/model_zoo/{ => flux2}/flux2_dev.py (91%) rename lightx2v_train/lightx2v_train/model_zoo/{ => flux2}/flux2_klein.py (90%) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/__init__.py rename lightx2v_train/lightx2v_train/model_zoo/{ => longcat_image}/longcat_image.py (91%) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py rename lightx2v_train/lightx2v_train/model_zoo/{ => ltx}/ltx_t2av.py (92%) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/qwen_image/__init__.py rename lightx2v_train/lightx2v_train/model_zoo/{ => qwen_image}/qwen_image.py (93%) rename lightx2v_train/lightx2v_train/model_zoo/{ => qwen_image}/qwen_image_edit.py (96%) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_teacher_forcing_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py rename lightx2v_train/lightx2v_train/model_zoo/{ => wan}/lingbot_video.py (92%) rename lightx2v_train/lightx2v_train/model_zoo/{ => wan}/wan_fastwam.py (91%) rename lightx2v_train/lightx2v_train/model_zoo/{ => wan}/wan_t2v.py (90%) rename lightx2v_train/lightx2v_train/model_zoo/{ => wan}/wan_ti2v_5b.py (95%) create mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/ltx_trainer.py create mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/runtime.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/video_ar_trainer.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/video_trainer.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/flow.py create mode 100644 lightx2v_train/lightx2v_train/trainers/flow_matching.py create mode 100644 lightx2v_train/lightx2v_train/trainers/teacher_forcing.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/tf.py rename lightx2v_train/lightx2v_train/trainers/{fastwam.py => world_action.py} (95%) create mode 100644 lightx2v_train/lightx2v_train/utils/image_size_buckets.py 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/wan.yaml b/lightx2v_train/configs/infer/wan.yaml index a2cc480e0..a66d54aa6 100644 --- a/lightx2v_train/configs/infer/wan.yaml +++ b/lightx2v_train/configs/infer/wan.yaml @@ -38,7 +38,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 pin_memory: true shuffle: false 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 index ef31546f9..bbad49df0 100644 --- a/lightx2v_train/configs/infer/wan2_2_ti2v_5b_tf_ar.yaml +++ b/lightx2v_train/configs/infer/wan2_2_ti2v_5b_tf_ar.yaml @@ -49,7 +49,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 10 - batch_size: 1 num_workers: 0 pin_memory: true shuffle: false diff --git a/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml index 3280370a1..7641917b9 100644 --- a/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml +++ b/lightx2v_train/configs/train/dmd/flux2_dev_cdm_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_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..5e13abb71 100644 --- a/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/flux2_klein_dmd_lora.yaml @@ -6,13 +6,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 @@ -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 index 2b800df1e..cf3d969ab 100644 --- a/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml +++ b/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml @@ -7,13 +7,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 @@ -53,9 +54,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/longcat_image_dmd_lora.yaml b/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml index 6a1b65d8e..4fd658dbf 100644 --- a/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/longcat_image_dmd_lora.yaml @@ -7,13 +7,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 @@ -53,9 +54,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 index 77b9aecf5..95bdab321 100644 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml @@ -49,7 +49,6 @@ 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 @@ -69,7 +68,7 @@ scheduler: extra_one_step: true training: - method: ltx_t2av_ar_dmd + method: autoregressive_dmd max_train_iters: 1000 gradient_accumulation_iters: 4 gradient_checkpointing: 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 index 878fdae33..6bba4edf4 100644 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml +++ b/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml @@ -40,7 +40,6 @@ 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 @@ -53,7 +52,7 @@ scheduler: max_t: 1.0 training: - method: ltx_t2av_dmd + method: dmd max_train_iters: 1000 gradient_accumulation_iters: 4 gradient_checkpointing: true diff --git a/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml index 68cca7fd1..f268a6dfd 100644 --- a/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml +++ b/lightx2v_train/configs/train/dmd/qwen_image_cdm_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/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 index b1331a448..452109787 100644 --- 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 @@ -59,7 +59,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 @@ -69,7 +68,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 pin_memory: true shuffle: false @@ -88,7 +86,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_2_ti2v_5b_dmd.yaml b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml index 16a021c87..116e74395 100644 --- a/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml +++ b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd.yaml @@ -48,7 +48,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 @@ -58,7 +57,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -76,7 +74,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_2_ti2v_5b_dmd_multi.yaml b/lightx2v_train/configs/train/dmd/wan2_2_ti2v_5b_dmd_multi.yaml index 7571d4047..19f7c922b 100644 --- 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 @@ -49,7 +49,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 @@ -59,7 +58,6 @@ data: data_path: - /path/to/prompts/wan_t2v_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -77,7 +75,7 @@ scheduler: extra_one_step: true training: - method: video_dmd + method: dmd max_train_iters: 5000 gradient_accumulation_iters: 2 gradient_checkpointing: 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 index efc546dcc..44b094c23 100644 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml @@ -13,17 +13,16 @@ data: 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 + method: world_action 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 + gradient_accumulation_iters: 32 max_grad_norm: 1.0 # One checkpoint per approximately 2,238,036 / 128 optimizer steps. save_every_iters: 17485 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 index eb6388d8d..fd0db2675 100644 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml @@ -13,17 +13,16 @@ data: 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 + method: world_action 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 + gradient_accumulation_iters: 16 max_grad_norm: 1.0 # One checkpoint per approximately 174848 / 8 optimizer steps. save_every_iters: 17485 diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index 4f20d265e..24bed494d 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -12,17 +12,16 @@ data: - /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 + method: world_action 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 + gradient_accumulation_iters: 32 max_grad_norm: 1.0 save_every_iters: 2000 save_total_limit: 2 diff --git a/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml b/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml index 7dcab5247..8d32cbd91 100644 --- a/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml +++ b/lightx2v_train/configs/train/flow/flux2_dev_lora.yaml @@ -54,7 +54,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 diff --git a/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml b/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml index 24f86bf06..6375c2c12 100644 --- a/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml +++ b/lightx2v_train/configs/train/flow/flux2_klein_lora.yaml @@ -40,7 +40,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 diff --git a/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml b/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml index b651ba0a5..5ef54d074 100644 --- a/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml +++ b/lightx2v_train/configs/train/flow/flux2_klein_t2i.yaml @@ -54,7 +54,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_lora.yaml b/lightx2v_train/configs/train/flow/longcat_image_lora.yaml index 2998517d0..3831edb0a 100644 --- a/lightx2v_train/configs/train/flow/longcat_image_lora.yaml +++ b/lightx2v_train/configs/train/flow/longcat_image_lora.yaml @@ -43,7 +43,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/longcat_image_t2i.yaml b/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml index c8ec9196c..dd6e0f24c 100644 --- a/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml +++ b/lightx2v_train/configs/train/flow/longcat_image_t2i.yaml @@ -57,7 +57,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/ltx_t2av_full.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml index e1b07768f..2026108c4 100644 --- a/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml +++ b/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml @@ -30,7 +30,6 @@ data: 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 @@ -42,7 +41,7 @@ scheduler: logitnormal_uniform_prob: 0.1 training: - method: ltx_t2av_flow + method: flow_matching train_type: full max_train_iters: 3000 gradient_accumulation_iters: 2 diff --git a/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml index cb49f24e7..7858854b7 100644 --- a/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml +++ b/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml @@ -30,7 +30,6 @@ data: name: latent_dataset data_path: /path/to/ltx_t2av/.precomputed dataset_repeat: 1 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -42,7 +41,7 @@ scheduler: logitnormal_uniform_prob: 0.1 training: - method: ltx_t2av_flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 1 diff --git a/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml index 0e95bad41..9936a3d6b 100644 --- a/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml +++ b/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml @@ -30,7 +30,6 @@ data: 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 @@ -42,7 +41,7 @@ scheduler: logitnormal_uniform_prob: 0.1 training: - method: ltx_t2av_flow + method: flow_matching train_type: lora max_train_iters: 3000 gradient_accumulation_iters: 2 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 e02555932..1bc8e1193 100644 --- a/lightx2v_train/configs/train/flow/qwen_image_lora.yaml +++ b/lightx2v_train/configs/train/flow/qwen_image_lora.yaml @@ -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_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..a7962a71d 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 @@ -44,7 +44,6 @@ data: random_start: false skip_missing: true decode_retries: 3 - batch_size: 1 num_workers: 4 pin_memory: true shuffle: true @@ -53,7 +52,6 @@ data: data_path: - /path/to/wan_t2v_prompts.txt max_samples: 5 - batch_size: 1 num_workers: 0 shuffle: false @@ -77,7 +75,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 index 309f6fe7e..b7f89a6fa 100644 --- a/lightx2v_train/configs/train/sgmd/wan2_2_ti2v_5b_sgmd.yaml +++ b/lightx2v_train/configs/train/sgmd/wan2_2_ti2v_5b_sgmd.yaml @@ -48,7 +48,6 @@ data: name: prompt_dataset data_path: - data/prompts/wan22_train_prompts.txt - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -58,7 +57,6 @@ data: data_path: - data/prompts/wan22_val_prompts.txt max_samples: 1 - batch_size: 1 num_workers: 0 shuffle: false @@ -76,7 +74,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/tf/ltx_t2av_teacher_forcing_lmdb.yaml b/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml index 8513c89b4..644522f65 100644 --- a/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml +++ b/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml @@ -35,7 +35,6 @@ data: 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 @@ -55,7 +54,7 @@ scheduler: extra_one_step: true training: - method: ltx_t2av_teacher_forcing + method: teacher_forcing train_type: full max_train_iters: 3000 gradient_accumulation_iters: 4 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..bd9f40d7a 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 @@ -33,7 +33,6 @@ data: name: latent_dataset data_path: /path/to/datasets/causal_forcing_data/dataset/clean_data/ dataset_repeat: 1 - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -43,7 +42,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/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..328960cbc 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 @@ -33,7 +33,6 @@ data: name: latent_dataset data_path: /path/to/datasets/causal_forcing_data/dataset/clean_data dataset_repeat: 1 - batch_size: 1 num_workers: 8 pin_memory: true shuffle: true @@ -43,7 +42,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/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 index 78a4a8abc..c9258e6d5 100644 --- 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 @@ -44,7 +44,6 @@ data: 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 @@ -54,7 +53,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/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 index dc88c14f1..2b4b63364 100644 --- 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 @@ -43,7 +43,6 @@ data: 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 @@ -53,7 +52,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/data_process/README.md b/lightx2v_train/data_process/README.md index 5464f1e01..200169bca 100644 --- a/lightx2v_train/data_process/README.md +++ b/lightx2v_train/data_process/README.md @@ -12,6 +12,7 @@ data_process/ | 训练方式 | 数据集 | 输入 | | --- | --- | --- | +| 图像 DMD / CDM | `prompt_dataset` | prompt、目标高度和目标宽度 | | Flow | `video_dataset` | 视频路径和 caption 元数据，训练时在线编码 | | Flow | `latent_dataset` | 预先生成的视频、音频和文本缓存 | | DMD / AR-DMD | `prompt_dataset` | TXT prompt，训练时在线编码 | @@ -52,23 +53,82 @@ data: width: 832 num_frames: 81 frame_rate: 24 - batch_size: 1 ``` 视频路径可以是绝对路径，也可以是相对元数据文件或 `video_root` 的路径。 ## prompt_dataset -`prompt_dataset` 读取 TXT 或 LIST 文件，每个非空行表示一个 prompt。 +`prompt_dataset` 支持 TXT、LIST、JSON 和 JSONL；TXT/LIST 中每个非空行表示一个 prompt。 ```yaml data: train: name: prompt_dataset data_path: /path/to/prompts/train.txt - batch_size: 1 ``` +图像 DMD / CDM 使用 JSON 或 JSONL。每条记录必须显式指定目标尺寸， +不需要 `target_image`： + +```json +{"prompt":"A cat walking through snow.","target_height":1024,"target_width":1024} +``` + +多尺寸训练时，每条记录的 `(target_height, target_width)` 表示它所属的 +精确尺寸桶。同一 batch 的样本会由 `bucket_by_size` 按尺寸分组： + +```jsonl +{"prompt":"A square composition.","target_height":1024,"target_width":1024} +{"prompt":"A wide landscape.","target_height":768,"target_width":1344} +{"prompt":"A tall portrait.","target_height":1344,"target_width":768} +``` + +```yaml +data: + train: + name: prompt_dataset + data_path: /path/to/image_dmd/train.jsonl + bucket_by_size: true + shuffle: true + +training: + method: dmd + dmd: + # 不配置采样比例：保持数据集中各尺寸桶的数量分布。 + image_sizes: + - {value: [1024, 1024]} + - {value: [16, 768, 1344]} + - {value: [16, 1344, 768]} +``` + +如果需要显式控制尺寸桶的 global-batch 采样比例，为每个条目添加 +`ratio`： + +```yaml +training: + dmd: + image_sizes: + - {value: [1024, 1024], ratio: 5} + - {value: [16, 768, 1344], ratio: 3} + - {value: [16, 1344, 768], ratio: 2} +``` + +`ratio` 是相对权重，会自动归一化；上例三个桶的期望采样占比为 +`50% / 30% / 20%`。权重在分布式训练步级别生效，所有分布式 rank +在同一训练步使用同一尺寸。`value` 只支持 `[height, width]` 或 +`[prefix, height, width]`，图像 DMD 使用最后两个值作为像素高宽。一份配置必须 +全部带 `ratio` 或全部不带，不能混用。旧的裸数组格式（例如 +`- [1024, 1024]`）不再支持。 + +如果配置了 `training.dmd.image_sizes`，每条样本的尺寸必须属于该列表； +如果不配置，则接受能被模型图像尺寸倍数整除的任意正尺寸（当前图像模型为 +VAE 空间压缩倍率的 2 倍）。每个样本都必须 +同时提供 `target_height` 和 `target_width`。物理 batch 固定为 1； +`bucket_by_size: true` 在分布式训练时保证同一步所有 rank 使用相同尺寸。 +默认 `drop_last: false` 会在每个尺寸桶内重复少量样本，使样本数能被数据并行 +rank 数整除；设置 `drop_last: true` 则丢弃不足一个分布式训练步的尾部样本。 + ## latent_dataset `latent_dataset` 读取 JSON、JSONL、CSV 缓存元数据、包含 `metadata.jsonl` 的目录或 LMDB。缓存元数据支持以下字段： @@ -98,12 +158,11 @@ 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`。 +常用 DataLoader 配置包括 `dataset_repeat`、`max_samples`、`num_workers`、`shuffle`、`drop_last` 和 `pin_memory`。物理 batch 固定为 1，数据配置中不存在 `batch_size` 选项。 ## 预处理输入元数据 @@ -216,7 +275,6 @@ data: width: 832 num_frames: 81 frame_rate: 24 - batch_size: 1 num_workers: 4 shuffle: true ``` @@ -228,7 +286,6 @@ data: train: name: latent_dataset data_path: /path/to/datasets/latent_cache/metadata.jsonl - batch_size: 1 num_workers: 4 shuffle: true ``` @@ -244,7 +301,6 @@ data: train: name: prompt_dataset data_path: /path/to/prompts/train.txt - batch_size: 1 num_workers: 4 shuffle: true drop_last: true @@ -257,7 +313,6 @@ 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 @@ -274,7 +329,6 @@ 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 @@ -287,7 +341,6 @@ data: train: name: latent_dataset data_path: /path/to/datasets/tf_cache.lmdb - batch_size: 1 num_workers: 4 shuffle: true drop_last: true diff --git a/lightx2v_train/lightx2v_train/data/libero/dataset.py b/lightx2v_train/lightx2v_train/data/libero/dataset.py index 5e20b265e..a6549b4c4 100644 --- a/lightx2v_train/lightx2v_train/data/libero/dataset.py +++ b/lightx2v_train/lightx2v_train/data/libero/dataset.py @@ -213,7 +213,7 @@ def build_libero_fastwam_dataset(data_config, train_or_val="train"): shuffle = False return DataLoader( dataset, - batch_size=int(data_config.get("batch_size", 1)), + batch_size=1, shuffle=shuffle, sampler=sampler, num_workers=int(data_config.get("num_workers", 4)), diff --git a/lightx2v_train/lightx2v_train/data/libero/preparation.py b/lightx2v_train/lightx2v_train/data/libero/preparation.py index 38524228c..f0db2db98 100644 --- a/lightx2v_train/lightx2v_train/data/libero/preparation.py +++ b/lightx2v_train/lightx2v_train/data/libero/preparation.py @@ -218,10 +218,10 @@ def precompute_text_embeddings(model_path, cache_dir, context_len, prompts): tokenizer_path=str(tokenizer_path), ) encoder.model.eval() - batch_size = 16 + encoding_chunk_size = 16 with torch.no_grad(): - for start in range(0, len(missing_prompts), batch_size): - batch_prompts = missing_prompts[start : start + batch_size] + for start in range(0, len(missing_prompts), encoding_chunk_size): + batch_prompts = missing_prompts[start : start + encoding_chunk_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) diff --git a/lightx2v_train/lightx2v_train/data/utils.py b/lightx2v_train/lightx2v_train/data/utils.py index 265b4b5d5..0cca7eb9c 100644 --- a/lightx2v_train/lightx2v_train/data/utils.py +++ b/lightx2v_train/lightx2v_train/data/utils.py @@ -14,6 +14,12 @@ _BILINEAR = getattr(Image, "Resampling", Image).BILINEAR +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 [] diff --git a/lightx2v_train/lightx2v_train/data/video_dataset.py b/lightx2v_train/lightx2v_train/data/video_dataset.py index 916915725..2ac6d5b09 100644 --- a/lightx2v_train/lightx2v_train/data/video_dataset.py +++ b/lightx2v_train/lightx2v_train/data/video_dataset.py @@ -1,6 +1,7 @@ import io import json import random +from collections import defaultdict from pathlib import Path import numpy as np @@ -18,6 +19,7 @@ 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"} @@ -278,11 +280,24 @@ def _load_prompts(self): negative_prompt = json.dumps(negative_prompt, ensure_ascii=False, separators=(",", ":")) else: negative_prompt = str(negative_prompt).strip() + 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 samples.append( { "prompt": prompt, "negative_prompt": negative_prompt, - "meta": {"prompt_path": str(path), "row_index": row_index}, + "meta": meta, } ) if self.max_samples is not None and len(samples) >= self.max_samples: @@ -306,6 +321,169 @@ def __getitem__(self, index): def __len__(self): return len(self.samples) * self.dataset_repeat + 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 int(height), int(width) + + +class SizeBucketSampler(torch.utils.data.Sampler): + """Keep all data-parallel ranks on one exact size bucket per step.""" + + 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 self.num_samples + class LatentDataset(torch.utils.data.Dataset): def __init__( @@ -519,7 +697,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, @@ -531,7 +725,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), 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..1d4b132b1 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -0,0 +1,49 @@ +from .base import ( + BoundCapability, + CapabilityNotSupportedError, + CapabilityProvider, + CapabilityRegistry, + ModelCapability, +) +from .training import ( + AutoregressiveDistillationCapability, + AutoregressiveRolloutContext, + CheckpointCapability, + ConsistencyCapability, + DistillationCapability, + DopsdCapability, + DopsdLossResult, + DopsdStepContext, + FlowMatchingSFTCapability, + LossResult, + ParallelCapability, + SFTStepContext, + TeacherForcingCapability, + TeacherForcingStepContext, + TrainableModelCapability, + WorldActionTrainingCapability, +) + +__all__ = [ + "AutoregressiveDistillationCapability", + "AutoregressiveRolloutContext", + "BoundCapability", + "CapabilityNotSupportedError", + "CapabilityProvider", + "CapabilityRegistry", + "CheckpointCapability", + "ConsistencyCapability", + "DistillationCapability", + "DopsdCapability", + "DopsdLossResult", + "DopsdStepContext", + "FlowMatchingSFTCapability", + "LossResult", + "ModelCapability", + "ParallelCapability", + "SFTStepContext", + "TeacherForcingCapability", + "TeacherForcingStepContext", + "TrainableModelCapability", + "WorldActionTrainingCapability", +] 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..565eb2eb3 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/base.py @@ -0,0 +1,94 @@ +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 " + f"{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..9882f3379 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +from abc import abstractmethod +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping + +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) + + +@dataclass(frozen=True) +class SFTStepContext: + noise_scheduler: Any + running_dtype: torch.dtype + broadcast: Callable[[Any], Any] + + +class FlowMatchingSFTCapability(ModelCapability): + @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): + @abstractmethod + def compute_loss( + self, + batch: Mapping[str, Any], + context: TeacherForcingStepContext, + ) -> LossResult: + """Compute a differentiable teacher-forcing loss.""" + + +class ConsistencyCapability(ModelCapability): + @abstractmethod + def configure(self, features) -> None: + pass + + @abstractmethod + def restore_trainable_auxiliary(self) -> None: + pass + + @abstractmethod + def encode_latent(self, batch): + pass + + @abstractmethod + def encode_condition(self, batch): + pass + + @abstractmethod + def predict(self, request, path): + pass + + @abstractmethod + def predict_log_variance(self, time): + pass + + @abstractmethod + def set_frozen(self, training: bool = False) -> None: + pass + + @abstractmethod + def denoiser(self): + pass + + +class DistillationCapability(ModelCapability): + """Model operations required by distribution-matching algorithms.""" + + @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 AutoregressiveDistillationCapability(ModelCapability): + @abstractmethod + def rollout( + self, + condition, + latent_shape, + initial_latents, + context: AutoregressiveRolloutContext, + ): + """Run a differentiable chunk-wise autoregressive rollout.""" + + +@dataclass(frozen=True) +class DopsdStepContext: + scheduler: Any + num_training_steps: int + running_dtype: torch.dtype + student_adapter: str + teacher_adapter: str + teacher_prompts: Callable[[Any], list[str]] + step_loss_weights: Callable[[int], list[float]] + collect_trajectory: bool = False + + +@dataclass(frozen=True) +class DopsdLossResult: + loss: Tensor + student_trajectory: tuple[Tensor, ...] = () + teacher_trajectory: tuple[Tensor, ...] = () + latent_ids: Tensor | None = None + height: int | None = None + width: int | None = None + + +class DopsdCapability(ModelCapability): + @property + @abstractmethod + def device(self): + pass + + @abstractmethod + def configure_adapters( + self, + rank, + alpha, + target_modules, + student_adapter, + teacher_adapter, + initialize_teacher, + ) -> None: + pass + + @abstractmethod + def parameters(self): + pass + + @abstractmethod + def compute_loss( + self, + batch, + context: DopsdStepContext, + ) -> DopsdLossResult: + pass + + @abstractmethod + def ema_update(self, student_adapter, teacher_adapter, decay) -> None: + pass + + @abstractmethod + def decode_trajectory(self, trajectory, latent_ids): + pass + + @abstractmethod + def set_training(self, student_adapter, teacher_adapter) -> None: + pass + + @abstractmethod + def set_eval(self) -> None: + pass + + @abstractmethod + def set_active_adapter(self, adapter_name) -> None: + pass + + @abstractmethod + def encode_prompt(self, prompts): + pass + + @abstractmethod + def prepare_reference(self, image): + pass + + @abstractmethod + def initial_latents(self, height, width, generator=None): + pass + + @abstractmethod + def predict_velocity( + self, + latents, + time, + condition, + latent_ids, + adapter_name, + **kwargs, + ): + pass + + @abstractmethod + def load_adapter(self, path, adapter_name, weights_subdir=None) -> None: + pass + + @abstractmethod + def save_adapter(self, path, adapter_name, weights_subdir=None) -> None: + pass + + @abstractmethod + def copy_adapter(self, source_adapter, target_adapter) -> None: + pass + + +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 + + +class WorldActionTrainingCapability(ModelCapability): + @abstractmethod + def configure(self) -> None: + pass + + @abstractmethod + def parameters(self): + pass + + @abstractmethod + def module(self): + pass + + @abstractmethod + def compute_loss(self, batch: Mapping[str, Any], module=None) -> LossResult: + pass + + @abstractmethod + def evaluation_loss(self, batch: Mapping[str, Any]) -> LossResult: + pass + + @abstractmethod + def load_checkpoint(self, path) -> None: + pass + + @abstractmethod + def save_checkpoint(self, path, step=None) -> None: + pass diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 74af37547..a37dea7e2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -3,19 +3,35 @@ 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"), - "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"), + "Flux2KleinModel": (".flux2.flux2_klein", "Flux2KleinModel"), + "LingBotVideoModel": (".wan.lingbot_video", "LingBotVideoModel"), + "LongCatImageModel": (".longcat_image.longcat_image", "LongCatImageModel"), + "LTX2T2AVModel": (".ltx.ltx_t2av", "LTX2T2AVModel"), + "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), + "QwenImageEditModel": (".qwen_image.qwen_image_edit", "QwenImageEditModel"), + "WanFastWAMModel": (".wan.wan_fastwam", "WanFastWAMModel"), + "WanT2VModel": (".wan.wan_t2v", "WanT2VModel"), + "WanTI2V5BModel": (".wan.wan_ti2v_5b", "WanTI2V5BModel"), } +def build_loaded_model( + config, + *, + transformer_only=False, + reference_model=None, +): + """Build a model wrapper, load its components, then publish capabilities.""" + model = build_model(config) + model.load_components( + transformer_only=transformer_only, + reference_model=reference_model, + ) + model.ensure_capabilities() + return model + + def __getattr__(name): target = _LAZY_EXPORTS.get(name) if target is None: @@ -28,6 +44,7 @@ def __getattr__(name): __all__ = [ "build_model", + "build_loaded_model", "QwenImageModel", "QwenImageEditModel", "LongCatImageModel", diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 71b94c1d2..5e8dec550 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -12,24 +12,47 @@ 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): 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 + 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): raise NotImplementedError - def dmd_latent_shape(self, batch_size, height, width): - raise NotImplementedError(f"{self.__class__.__name__} must define dmd_latent_shape().") - def denoiser_module(self): raise NotImplementedError(f"{self.__class__.__name__} must define denoiser_module().") @@ -210,7 +233,7 @@ def predict_denoiser_output(self, noisy_latent, timestep_or_sigma, condition, ** def prepare_infer_latents(self, height, width, generator=None): raise NotImplementedError - def dmd_latent_shape(self, batch_size, height, width): + def dmd_latent_shape(self, height, width): raise NotImplementedError(f"{self.__class__.__name__} must define dmd_latent_shape().") def cfg_on_denoiser_output(self): 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..e3f56311f --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py @@ -0,0 +1 @@ +"""Capability adapters shared by multiple model families.""" 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..99a2a2b9a --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -0,0 +1,492 @@ +from __future__ import annotations + +import os +from typing import Any, Mapping + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_capabilities import ( + BoundCapability, + CheckpointCapability, + ConsistencyCapability, + DistillationCapability, + 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 + + +def _require_single_prompt(prompt): + if isinstance(prompt, str): + return True + prompts = list(prompt) + if len(prompts) != 1: + raise ValueError( + "Training requires exactly one prompt per rank; " + f"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( + "Training requires exactly one negative prompt per rank; " + f"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 " + f"than 1 are not supported, got {shape}." + ) + return value + + +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): + def save_weights(self, save_dir, train_type) -> None: + if train_type == "lora": + self.model.save_lora_weights(save_dir) + 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_consistency_auxiliary_weights(save_dir) + 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 compute_loss( + self, + batch: Mapping[str, Any], + context: SFTStepContext, + ) -> LossResult: + scheduler = context.noise_scheduler + broadcast = context.broadcast + + with torch.no_grad(): + latent = _require_singleton_tensor( + broadcast(self.model.encode_to_latent(batch)), + "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 = broadcast(self.model.encode_condition(batch)) + + 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 GenericConsistencyCapability(BoundCapability, ConsistencyCapability): + def configure(self, features) -> None: + self.model.configure_consistency_model(features) + + def restore_trainable_auxiliary(self) -> None: + self.model.set_consistency_modules_trainable() + + def encode_latent(self, batch): + return _require_singleton_tensor( + self.model.encode_to_latent(batch), + "Consistency latent", + ) + + def encode_condition(self, batch): + return self.model.encode_condition(batch) + + 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): + return self.model.predict_consistency_log_variance(time) + + 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 GenericDistillationCapability(BoundCapability, DistillationCapability): + """Distribution-matching operations shared by image flow models.""" + + @property + def device(self): + return self.model.device + + @property + def default_negative_prompt(self): + return None + + @property + def default_lora_target_modules(self): + return None + + 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( + "Image DMD sample size " + f"{height}x{width} is not in training.dmd.image_sizes: " + f"[{configured}]." + ) + + spatial_scale = int(self.model.vae_scale_factor) + size_multiple = spatial_scale * 2 + if height % size_multiple or width % size_multiple: + raise ValueError( + "Image DMD target_height and target_width must be divisible by " + f"the model image size multiple {size_multiple}, got " + f"{height}x{width}." + ) + return self.model.dmd_latent_shape(height, width) + + @staticmethod + def _target_dimension(meta, key): + if key not in meta: + raise KeyError( + "Image DMD prompt-only samples require meta.target_height and " + f"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) + 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.model.cfg_on_denoiser_output(): + denoiser_input = self.model.prepare_denoiser_input(latents) + 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 = batch["inputs"].get("latents") + if latent is None: + latent = self.model.encode_to_latent(batch) + latent = latent.to(device=self.device, dtype=dtype) + if latent.ndim == 4: + 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/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..175635986 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -0,0 +1 @@ +"""Capabilities implemented specifically by Flux2 models.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py new file mode 100644 index 000000000..82f4941b0 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py @@ -0,0 +1,201 @@ +"""DOPSD capability for Flux2 models.""" + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_capabilities import ( + BoundCapability, + DopsdCapability, + DopsdLossResult, + DopsdStepContext, +) + + +class Flux2DopsdCapability(BoundCapability, DopsdCapability): + @property + def device(self): + return self.model.device + + def configure_adapters( + self, + rank, + alpha, + target_modules, + student_adapter, + teacher_adapter, + initialize_teacher, + ) -> None: + self.model.add_dual_lora( + rank, + alpha, + target_modules, + student_adapter=student_adapter, + teacher_adapter=teacher_adapter, + init_teacher_from_student=initialize_teacher, + ) + self.set_training(student_adapter, teacher_adapter) + + def parameters(self): + return self.model.trainable_parameters() + + def compute_loss( + self, + batch, + context: DopsdStepContext, + ) -> DopsdLossResult: + image = batch["inputs"].get("target_image") + if image is None: + raise ValueError("D-OPSD requires inputs.target_image.") + image = image.to( + device=self.device, + dtype=context.running_dtype, + ) + if image.shape[0] != 1: + raise ValueError("D-OPSD only supports physical batch size 1.") + _, _, height, width = image.shape + latent_hw = (height // 16, width // 16) + timestep_scale = float(context.scheduler.num_train_timesteps) + + with torch.no_grad(): + student_condition = self.model.encode_condition(batch) + teacher_condition = self.model.encode_prompt_text(context.teacher_prompts(batch["conditioning"]["prompt"])) + reference_latents, reference_ids = self.model.prepare_reference_image_latents(image) + initial_latents, latent_ids = self.model.prepare_dopsd_initial_latents( + height, + width, + ) + context.scheduler.set_timesteps( + context.num_training_steps, + latent_hw=latent_hw, + ) + timesteps = context.scheduler.infer_timesteps + + latents = initial_latents + total_loss = 0.0 + weights = context.step_loss_weights(len(timesteps)) + student_trajectory = [] + teacher_trajectory = [] + for step_index, timestep in enumerate(timesteps): + time = (timestep.reshape(1) / timestep_scale).to( + device=self.device, + dtype=context.running_dtype, + ) + if step_index + 1 < len(timesteps): + next_time = timesteps[step_index + 1].reshape(1) + next_time = next_time / timestep_scale + else: + next_time = torch.zeros_like(time) + next_time = next_time.to( + device=self.device, + dtype=context.running_dtype, + ) + latents = latents.detach().requires_grad_(True) + with torch.no_grad(): + teacher_velocity = self.model.predict_velocity( + latents, + time, + teacher_condition, + latent_ids, + context.teacher_adapter, + teacher_image_latents=reference_latents, + teacher_image_latent_ids=reference_ids, + ) + teacher_x0 = latents - time.reshape(1, 1, 1) * teacher_velocity + student_velocity = self.model.predict_velocity( + latents, + time, + student_condition, + latent_ids, + context.student_adapter, + ) + student_x0 = latents - time.reshape(1, 1, 1) * student_velocity + latents = latents + student_velocity * (next_time - time).reshape(1, 1, 1) + total_loss = total_loss + weights[step_index] * F.mse_loss( + student_x0, + teacher_x0.detach(), + ) + if context.collect_trajectory: + student_trajectory.append(student_x0.detach()) + teacher_trajectory.append(teacher_x0.detach()) + + return DopsdLossResult( + loss=total_loss / sum(weights), + student_trajectory=tuple(student_trajectory), + teacher_trajectory=tuple(teacher_trajectory), + latent_ids=latent_ids if context.collect_trajectory else None, + height=height if context.collect_trajectory else None, + width=width if context.collect_trajectory else None, + ) + + def ema_update(self, student_adapter, teacher_adapter, decay) -> None: + self.model.ema_update_lora_adapter( + student_adapter, + teacher_adapter, + decay, + ) + + def decode_trajectory(self, trajectory, latent_ids): + return self.model.decode_packed_x0_to_images(trajectory, latent_ids) + + def set_training(self, student_adapter, teacher_adapter) -> None: + self.model.set_dual_lora_trainable( + student_adapter, + teacher_adapter, + ) + + def set_eval(self) -> None: + self.model.set_denoiser_eval() + + def set_active_adapter(self, adapter_name) -> None: + self.model.set_active_adapter(adapter_name) + + def encode_prompt(self, prompts): + return self.model.encode_prompt_text(prompts) + + def prepare_reference(self, image): + return self.model.prepare_reference_image_latents(image) + + def initial_latents(self, height, width, generator=None): + return self.model.prepare_dopsd_initial_latents( + height, + width, + generator=generator, + ) + + def predict_velocity( + self, + latents, + time, + condition, + latent_ids, + adapter_name, + **kwargs, + ): + return self.model.predict_velocity( + latents, + time, + condition, + latent_ids, + adapter_name, + **kwargs, + ) + + def load_adapter(self, path, adapter_name, weights_subdir=None) -> None: + self.model.load_lora_weights_for_resume( + path, + adapter_name=adapter_name, + weights_subdir=weights_subdir, + ) + + def save_adapter(self, path, adapter_name, weights_subdir=None) -> None: + self.model.save_lora_weights( + path, + adapter_name=adapter_name, + weights_subdir=weights_subdir, + ) + + def copy_adapter(self, source_adapter, target_adapter) -> None: + self.model.copy_lora_adapter_weights( + source_adapter, + target_adapter, + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py similarity index 91% rename from lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py rename to lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index c140f1ad4..2bd5da334 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -4,9 +4,11 @@ from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER -from .base import BaseModel +from ..base import BaseModel @dataclass @@ -21,6 +23,17 @@ class Flux2DevDenoiserInput: class Flux2DevModel(BaseModel): pipeline_cls = Flux2Pipeline + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + GenericDistillationCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): model_config = self.config["model"] model_path = model_config["pretrained_model_name_or_path"] @@ -146,10 +159,10 @@ def prepare_infer_latents(self, height, width, generator=None): 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): + def dmd_latent_shape(self, 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) + return (1, int(self.transformer.config.in_channels), latent_h // 2, latent_w // 2) def cfg_on_denoiser_output(self): return True diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py similarity index 90% rename from lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py rename to lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index c638c476b..09137d679 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -4,9 +4,14 @@ from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from lightx2v_train.model_capabilities import DistillationCapability, DopsdCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.flux2.capability_adapters.flux2_dopsd_capability import ( + Flux2DopsdCapability, +) from lightx2v_train.utils.registry import MODEL_REGISTER -from .base import BaseModel +from ..base import BaseModel @dataclass @@ -21,6 +26,21 @@ class Flux2KleinDenoiserInput: class Flux2KleinModel(BaseModel): pipeline_cls = Flux2KleinPipeline + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + GenericDistillationCapability(self), + ) + self.capabilities.register( + DopsdCapability, + Flux2DopsdCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): if transformer_only: if reference_model is not None: @@ -106,11 +126,11 @@ def encode_prompt_condition(self, prompt): def encode_prompt_text(self, prompt): return self.encode_prompt_condition(prompt) - def dmd_latent_shape(self, batch_size, height, width): + def dmd_latent_shape(self, 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, + 1, self.transformer.config.in_channels, latent_h // 2, latent_w // 2, @@ -195,11 +215,11 @@ def prepare_reference_image_latents(self, images): 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): + def prepare_dopsd_initial_latents(self, height, width, generator=None): pipeline = self.assemble_pipeline() num_latents_channels = self.transformer.config.in_channels // 4 return pipeline.prepare_latents( - batch_size=batch_size, + batch_size=1, num_latents_channels=num_latents_channels, height=height, width=width, 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.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py similarity index 91% rename from lightx2v_train/lightx2v_train/model_zoo/longcat_image.py rename to lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py index 68cf5b990..d416206a7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -6,9 +6,11 @@ from diffusers.models.transformers import LongCatImageTransformer2DModel from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids +from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER -from .base import BaseModel +from ..base import BaseModel @dataclass @@ -23,6 +25,17 @@ class LongCatImageDenoiserInput: class LongCatImageModel(BaseModel): pipeline_cls = LongCatImagePipeline + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + GenericDistillationCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): if transformer_only: if reference_model is not None: @@ -148,10 +161,10 @@ def prepare_infer_latents(self, height, width, generator=None): 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): + def dmd_latent_shape(self, 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) + return (1, int(self.vae.config.latent_channels), latent_h, latent_w) def cfg_on_denoiser_output(self): return True diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py new file mode 100644 index 000000000..c58c1964d --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py @@ -0,0 +1 @@ +"""LTX model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py new file mode 100644 index 000000000..e7c260825 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py @@ -0,0 +1 @@ +"""Capabilities implemented specifically by LTX models.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py new file mode 100644 index 000000000..7b08f010d --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py @@ -0,0 +1,21 @@ +"""Shared latent value objects for LTX training capabilities.""" + +from dataclasses import dataclass + +from torch import Tensor + + +@dataclass(frozen=True) +class LTXJointLatentShape: + video_latent: tuple[int, ...] + audio_latent: tuple[int, ...] + video_tokens: tuple[int, ...] + audio_tokens: tuple[int, ...] + fps: float + + +@dataclass(frozen=True) +class LTXJointLatents: + video: Tensor + audio: Tensor + shape: LTXJointLatentShape diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py new file mode 100644 index 000000000..2396fac06 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py @@ -0,0 +1,387 @@ +"""Autoregressive distillation capability for LTX models.""" + +from __future__ import annotations + +import math + +import torch + +from lightx2v_train.model_capabilities import ( + AutoregressiveDistillationCapability, + AutoregressiveRolloutContext, + BoundCapability, +) +from lightx2v_train.model_zoo.native.ltx2 import Modality + +from .common import LTXJointLatents +from .ltx_distillation_capability import LTXDistillationCapability + + +class LTXAutoregressiveDistillationCapability( + BoundCapability, + AutoregressiveDistillationCapability, +): + """Cached chunk-wise rollout for LTX joint video/audio states.""" + + def __init__(self, model) -> None: + super().__init__(model) + self.distillation = LTXDistillationCapability(model) + + def rollout( + self, + condition, + latent_shape, + initial_latents, + context: AutoregressiveRolloutContext, + ): + latents = initial_latents + if latents.video.shape[0] != 1 or latents.audio.shape[0] != 1: + raise ValueError("LTX autoregressive DMD only supports physical batch size 1.") + _, _, num_frames, latent_height, latent_width = latent_shape.video_latent + frame_sequence_length = int(latent_height) * int(latent_width) + if latents.video.shape[1] != num_frames * frame_sequence_length: + if latents.video.shape[1] % num_frames: + raise ValueError("LTX video token count is not divisible by latent frames.") + frame_sequence_length = latents.video.shape[1] // num_frames + if num_frames % context.frames_per_chunk: + raise ValueError(f"LTX latent frames ({num_frames}) must be divisible by frames_per_chunk ({context.frames_per_chunk}).") + + transformer = self._unwrap_transformer() + if transformer is None: + raise RuntimeError("LTX autoregressive rollout requires transformer KV-cache inference support.") + transformer.train() + num_blocks = num_frames // context.frames_per_chunk + exit_indices = self._sample_exit_indices( + 1 if context.same_step_across_blocks else num_blocks, + len(context.denoising_steps), + latents.video.device, + ) + video_positions = self.distillation._video_token_positions(latent_shape) + audio_positions = self.distillation._audio_token_positions(latent_shape) + cache = self._new_caches( + dtype=latents.video.dtype, + device=latents.video.device, + video_total_tokens=latents.video.shape[1], + audio_total_tokens=latents.audio.shape[1], + video_frame_sequence_length=frame_sequence_length, + num_frames=num_frames, + frames_per_chunk=context.frames_per_chunk, + ) + + video_chunks = [] + audio_chunks = [] + for block_index in range(num_blocks): + frame_start = block_index * context.frames_per_chunk + frame_end = frame_start + context.frames_per_chunk + video_start = frame_start * frame_sequence_length + video_end = frame_end * frame_sequence_length + audio_start, audio_end = self._audio_range( + block_index, + num_blocks, + latents.audio.shape[1], + ) + block = LTXJointLatents( + latents.video[:, video_start:video_end], + latents.audio[:, audio_start:audio_end], + latent_shape, + ) + positions = ( + video_positions[:, :, video_start:video_end], + audio_positions[:, :, audio_start:audio_end], + ) + exit_index = int(exit_indices[0] if context.same_step_across_blocks else exit_indices[block_index]) + + x0 = None + for step_index in range(exit_index + 1): + sigma = context.trajectory_scheduler.sigma_at( + step_index, + device=self.model.device, + dtype=context.running_dtype, + ) + enable_grad = context.grad_enabled and step_index == exit_index + grad_context = torch.enable_grad if enable_grad else torch.no_grad + with grad_context(): + velocity = self._predict_chunk( + block, + sigma, + condition, + positions, + cache, + video_start, + audio_start, + ) + x0 = self.distillation.x0_from_velocity( + block, + velocity, + sigma, + ) + if step_index < exit_index: + next_sigma = context.trajectory_scheduler.sigma_at( + step_index + 1, + device=self.model.device, + dtype=context.running_dtype, + ) + with torch.no_grad(): + noise = self.distillation.random_noise_like( + x0, + x0.video.dtype, + lambda value: value, + ) + block = self.distillation.add_noise( + context.trajectory_scheduler, + x0, + noise, + next_sigma, + ) + + video_chunks.append(x0.video) + audio_chunks.append(x0.audio) + cache_latents = self.distillation.detach(x0) + cache_sigma = torch.full( + (1,), + context.context_noise, + device=self.model.device, + dtype=context.running_dtype, + ) + if context.context_noise > 0: + cache_latents = self.distillation.add_noise( + context.trajectory_scheduler, + cache_latents, + self.distillation.random_noise_like( + cache_latents, + cache_latents.video.dtype, + lambda value: value, + ), + cache_sigma, + ) + with torch.no_grad(): + self._predict_chunk( + cache_latents, + cache_sigma, + condition, + positions, + cache, + video_start, + audio_start, + ) + + generated = LTXJointLatents( + torch.cat(video_chunks, dim=1).to(context.running_dtype), + torch.cat(audio_chunks, dim=1).to(context.running_dtype), + latent_shape, + ) + return generated, exit_indices + + def _predict_chunk( + self, + latents, + sigma, + condition, + positions, + cache, + video_start, + audio_start, + ): + video_context, audio_context, context_mask = condition + video_positions, audio_positions = positions + if latents.video.shape[0] != 1 or latents.audio.shape[0] != 1: + raise ValueError("LTX autoregressive DMD only supports physical batch size 1.") + sigma = sigma.to( + device=self.model.device, + dtype=self.model.running_dtype, + ) + video = Modality( + enabled=True, + latent=latents.video.to( + device=self.model.device, + dtype=self.model.running_dtype, + ), + sigma=sigma, + timesteps=sigma.view(-1, 1) + .expand( + 1, + latents.video.shape[1], + ) + .clone(), + positions=video_positions.to(device=self.model.device), + context=video_context.to( + device=self.model.device, + dtype=self.model.running_dtype, + ), + context_mask=context_mask.to(device=self.model.device), + ) + audio = Modality( + enabled=True, + latent=latents.audio.to( + device=self.model.device, + dtype=self.model.running_dtype, + ), + sigma=sigma, + timesteps=sigma.view(-1, 1) + .expand( + 1, + latents.audio.shape[1], + ) + .clone(), + positions=audio_positions.to(device=self.model.device), + context=audio_context.to( + device=self.model.device, + dtype=self.model.running_dtype, + ), + context_mask=context_mask.to(device=self.model.device), + ) + with self.model.transformer_forward_context(): + video_prediction, audio_prediction = self.model.denoiser_module()( + video=video, + audio=audio, + perturbations=None, + kv_cache=cache, + video_current_start=video_start, + audio_current_start=audio_start, + ) + return LTXJointLatents( + video_prediction, + audio_prediction, + latents.shape, + ) + + def _new_caches( + self, + dtype, + device, + video_total_tokens, + audio_total_tokens, + video_frame_sequence_length, + num_frames, + frames_per_chunk, + ): + transformer = self._unwrap_transformer() + num_layers = len(transformer.transformer_blocks) + video_heads = int(transformer.num_attention_heads) + video_head_dim = int(transformer.inner_dim / video_heads) + audio_heads = int(transformer.audio_num_attention_heads) + audio_head_dim = int(transformer.audio_inner_dim / audio_heads) + local_attention_size = int(getattr(transformer, "local_attn_size", -1)) + sink_size = int(getattr(transformer, "sink_size", 0)) + video_chunk_tokens = frames_per_chunk * video_frame_sequence_length + if local_attention_size == -1: + video_cache_size = video_total_tokens + video_window = video_total_tokens + else: + video_window = max( + video_chunk_tokens, + local_attention_size * video_frame_sequence_length, + ) + video_cache_size = max( + video_window, + sink_size * video_frame_sequence_length + video_chunk_tokens, + ) + audio_tokens_per_frame = max( + 1, + math.ceil(audio_total_tokens / max(1, num_frames)), + ) + audio_chunk_tokens = max( + 1, + math.ceil(audio_total_tokens / (num_frames // frames_per_chunk)), + ) + if local_attention_size == -1: + audio_cache_size = audio_total_tokens + audio_window = audio_total_tokens + else: + audio_window = max( + audio_chunk_tokens, + local_attention_size * audio_tokens_per_frame, + ) + audio_cache_size = max( + audio_window, + sink_size * audio_tokens_per_frame + audio_chunk_tokens, + ) + return { + "video": [ + self._empty_cache( + video_cache_size, + video_heads, + video_head_dim, + dtype, + device, + video_window, + sink_size * video_frame_sequence_length, + ) + for _ in range(num_layers) + ], + "audio": [ + self._empty_cache( + audio_cache_size, + audio_heads, + audio_head_dim, + dtype, + device, + audio_window, + sink_size * audio_tokens_per_frame, + ) + for _ in range(num_layers) + ], + } + + @staticmethod + def _empty_cache( + cache_size, + heads, + head_dim, + dtype, + device, + attention_window_size, + sink_tokens, + ): + return { + "k": torch.zeros( + (1, cache_size, heads, head_dim), + dtype=dtype, + device=device, + ), + "v": torch.zeros( + (1, cache_size, heads, head_dim), + dtype=dtype, + device=device, + ), + "global_end_index": torch.zeros( + 1, + dtype=torch.long, + device=device, + ), + "local_end_index": torch.zeros( + 1, + dtype=torch.long, + device=device, + ), + "attention_window_size": int(attention_window_size), + "sink_tokens": int(sink_tokens), + } + + @staticmethod + def _audio_range(block_index, num_blocks, total_tokens): + start = (block_index * total_tokens + num_blocks - 1) // num_blocks + end = ((block_index + 1) * total_tokens + num_blocks - 1) // num_blocks + return start, end + + @staticmethod + def _sample_exit_indices(count, num_steps, device): + if torch.distributed.is_available() and torch.distributed.is_initialized(): + if torch.distributed.get_rank() == 0: + indices = torch.randint(0, num_steps, (count,), device=device) + else: + indices = torch.empty(count, dtype=torch.long, device=device) + torch.distributed.broadcast(indices, src=0) + return indices.tolist() + return torch.randint(0, num_steps, (count,), device=device).tolist() + + def _unwrap_transformer(self): + transformer = self.model.denoiser_module() + for candidate in ( + transformer, + getattr(transformer, "module", None), + getattr(transformer, "_fsdp_wrapped_module", None), + ): + if candidate is not None and hasattr(candidate, "_forward_inference"): + return candidate + return None diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py new file mode 100644 index 000000000..f05f3e8f1 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py @@ -0,0 +1,440 @@ +"""Distribution-matching capability for LTX joint video/audio models.""" + +from __future__ import annotations + +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistillationCapability, + _negative_prompt, + _require_single_prompt, + _require_singleton_tensor, +) +from lightx2v_train.model_zoo.native.ltx2 import ( + AudioLatentShape, + AudioPatchifier, + Modality, + SpatioTemporalScaleFactors, + VideoLatentPatchifier, + VideoLatentShape, + get_pixel_coords, +) +from lightx2v_train.utils.constants import LTX2_NEGATIVE_PROMPT + +from .common import LTXJointLatentShape, LTXJointLatents + + +class LTXDistillationCapability(GenericDistillationCapability): + """Distribution-matching operations for joint LTX video/audio states.""" + + def __init__(self, model) -> None: + super().__init__(model) + config = model.config["training"].get("ltx2", {}) + dmd_config = model.config["training"].get("dmd", {}) + self.video_weight = float(config.get("video_loss_weight", 1.0)) + self.audio_weight = float(config.get("audio_loss_weight", 1.0)) + self.default_fps = float(config.get("default_fps", dmd_config.get("fps", 25.0))) + self.video_patchifier = VideoLatentPatchifier(patch_size=1) + self.audio_patchifier = AudioPatchifier(patch_size=1) + + @property + def default_negative_prompt(self): + return LTX2_NEGATIVE_PROMPT + + @property + def default_lora_target_modules(self): + return None + + def latent_shape( + self, + batch, + shape_config, + image_sizes, + broadcast, + ): + del image_sizes, broadcast + prompt = batch["conditioning"].get("prompt", "") + _require_single_prompt(prompt) + video_shape = shape_config.get("video_latent_shape") + if video_shape is None: + meta = batch["meta"] + height = int(meta["target_height"][0].item()) if "target_height" in meta else int(shape_config.get("height", 512)) + width = int(meta["target_width"][0].item()) if "target_width" in meta else int(shape_config.get("width", 768)) + frames = int(shape_config.get("num_frames", 241)) + video_shape = ( + 1, + int(shape_config.get("video_latent_channels", 128)), + (frames - 1) // int(shape_config.get("video_temporal_scale", 8)) + 1, + height // int(shape_config.get("video_spatial_scale", 32)), + width // int(shape_config.get("video_spatial_scale", 32)), + ) + else: + video_shape = self._singleton_shape( + video_shape, + "video_latent_shape", + ) + + audio_shape = shape_config.get("audio_latent_shape") + if audio_shape is None: + if "audio_num_frames" in shape_config: + audio_frames = int(shape_config["audio_num_frames"]) + else: + pixel_frames = int( + shape_config.get( + "num_frames", + (video_shape[2] - 1) * 8 + 1, + ) + ) + fps = float(shape_config.get("fps", self.default_fps)) + audio_frames = round(pixel_frames / fps * 25.0) + audio_shape = ( + 1, + int(shape_config.get("audio_latent_channels", 8)), + audio_frames, + int(shape_config.get("audio_mel_bins", 16)), + ) + else: + audio_shape = self._singleton_shape( + audio_shape, + "audio_latent_shape", + ) + return LTXJointLatentShape( + video_latent=video_shape, + audio_latent=audio_shape, + video_tokens=( + 1, + video_shape[2] * video_shape[3] * video_shape[4], + video_shape[1], + ), + audio_tokens=( + 1, + audio_shape[2], + audio_shape[1] * audio_shape[3], + ), + fps=float(shape_config.get("fps", self.default_fps)), + ) + + def encode_conditions( + self, + batch, + negative_prompt, + guidance_scale, + broadcast, + ): + conditioning = batch["conditioning"] + positive_value = conditioning.get("positive") + if positive_value is None: + return super().encode_conditions( + batch, + negative_prompt, + guidance_scale, + broadcast, + ) + with torch.no_grad(): + positive = self._cached_condition(positive_value) + if guidance_scale > 1: + negative_value = conditioning.get("negative") + if negative_value is not None: + negative = self._cached_condition(negative_value) + else: + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + 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 _cached_condition(self, value): + if isinstance(value, (list, tuple)) and len(value) == 3: + video, audio, mask = value + elif isinstance(value, dict) and { + "video_context", + "audio_context", + "context_mask", + }.issubset(value): + video = value["video_context"] + audio = value["audio_context"] + mask = value["context_mask"] + else: + return self.model.prepare_text_condition(value) + video = _require_singleton_tensor( + video.to(device=self.device, dtype=self.model.running_dtype), + "LTX video context", + ) + audio = _require_singleton_tensor( + audio.to(device=self.device, dtype=self.model.running_dtype), + "LTX audio context", + ) + mask = _require_singleton_tensor( + mask.to(device=self.device), + "LTX context mask", + ) + return video, audio, mask + + def predict_velocity(self, latents, sigma, condition): + video_context, audio_context, context_mask = condition + _require_singleton_tensor(latents.video, "LTX video latent") + _require_singleton_tensor(latents.audio, "LTX audio latent") + sigma = sigma.to(device=self.device, dtype=self.model.running_dtype) + video_timesteps = ( + sigma.view(-1, 1) + .expand( + 1, + latents.video.shape[1], + ) + .clone() + ) + audio_timesteps = ( + sigma.view(-1, 1) + .expand( + 1, + latents.audio.shape[1], + ) + .clone() + ) + video = Modality( + enabled=True, + latent=latents.video.to( + device=self.device, + dtype=self.model.running_dtype, + ), + sigma=sigma, + timesteps=video_timesteps, + positions=self._video_token_positions(latents.shape), + context=video_context.to( + device=self.device, + dtype=self.model.running_dtype, + ), + context_mask=context_mask.to(device=self.device), + ) + audio = Modality( + enabled=True, + latent=latents.audio.to( + device=self.device, + dtype=self.model.running_dtype, + ), + sigma=sigma, + timesteps=audio_timesteps, + positions=self._audio_token_positions(latents.shape), + context=audio_context.to( + device=self.device, + dtype=self.model.running_dtype, + ), + context_mask=context_mask.to(device=self.device), + ) + with self.model.transformer_forward_context(): + video_prediction, audio_prediction = self.model.denoiser_module()( + video=video, + audio=audio, + perturbations=None, + ) + return LTXJointLatents( + video_prediction, + audio_prediction, + latents.shape, + ) + + def predict_guided_velocity( + self, + latents, + sigma, + condition, + negative_condition, + guidance_scale, + cfg_norm, + ): + positive = self.predict_velocity(latents, sigma, condition) + if negative_condition is None: + return positive + negative = self.predict_velocity(latents, sigma, negative_condition) + return LTXJointLatents( + self._cfg( + positive.video, + negative.video, + guidance_scale, + cfg_norm, + ), + self._cfg( + positive.audio, + negative.audio, + guidance_scale, + cfg_norm, + ), + latents.shape, + ) + + def initial_latents(self, latent_shape, dtype, broadcast): + del broadcast + if latent_shape.video_tokens[0] != 1 or latent_shape.audio_tokens[0] != 1: + raise ValueError(f"LTX DMD latent shape must start with 1, got {latent_shape}.") + return LTXJointLatents( + torch.randn( + latent_shape.video_tokens, + device=self.device, + dtype=dtype, + ), + torch.randn( + latent_shape.audio_tokens, + device=self.device, + dtype=dtype, + ), + latent_shape, + ) + + @staticmethod + def latent_hw(latent_shape): + return None + + @staticmethod + def random_noise_like(latents, dtype, broadcast): + del broadcast + return LTXJointLatents( + torch.randn_like(latents.video, dtype=dtype), + torch.randn_like(latents.audio, dtype=dtype), + latents.shape, + ) + + @staticmethod + def add_noise(scheduler, latents, noise, sigma): + return LTXJointLatents( + scheduler.add_noise(latents.video, noise.video, sigma).to(latents.video.dtype), + scheduler.add_noise(latents.audio, noise.audio, sigma).to(latents.audio.dtype), + latents.shape, + ) + + @staticmethod + def training_target(latents, noise): + return LTXJointLatents( + noise.video - latents.video.float(), + noise.audio - latents.audio.float(), + latents.shape, + ) + + @staticmethod + def step(scheduler, velocity, step_index, sample): + video_next, video_x0 = scheduler.step_by_index( + velocity.video, + step_index, + sample.video, + ) + audio_next, audio_x0 = scheduler.step_by_index( + velocity.audio, + step_index, + sample.audio, + ) + return ( + LTXJointLatents(video_next, audio_next, sample.shape), + LTXJointLatents(video_x0, audio_x0, sample.shape), + ) + + @staticmethod + def x0_from_velocity(sample, velocity, sigma): + def project(value, flow): + expanded = sigma.reshape( + sigma.shape[0], + *([1] * (value.ndim - 1)), + ) + return value - expanded * flow + + return LTXJointLatents( + project(sample.video, velocity.video), + project(sample.audio, velocity.audio), + sample.shape, + ) + + def regression_loss(self, prediction, target): + video = torch.nn.functional.mse_loss( + prediction.video.float(), + target.video.float(), + ) + audio = torch.nn.functional.mse_loss( + prediction.audio.float(), + target.audio.float(), + ) + return self.video_weight * video + self.audio_weight * audio + + 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 LTXJointLatents( + value.video.detach(), + value.audio.detach(), + value.shape, + ) + + @staticmethod + def to_dtype(value, dtype): + return LTXJointLatents( + value.video.to(dtype=dtype), + value.audio.to(dtype=dtype), + value.shape, + ) + + def extract_real_latents(self, batch, dtype, broadcast): + raise ValueError("LTX joint DMD does not support real-data fake loss.") + + def _video_token_positions(self, shape): + _, channels, frames, height, width = shape.video_latent + latent_coords = self.video_patchifier.get_patch_grid_bounds( + output_shape=VideoLatentShape( + frames=frames, + height=height, + width=width, + batch=shape.video_latent[0], + channels=channels, + ), + device=self.device, + ) + positions = get_pixel_coords( + latent_coords=latent_coords, + scale_factors=SpatioTemporalScaleFactors.default(), + causal_fix=True, + ).float() + positions[:, 0, ...] /= shape.fps + return positions + + def _audio_token_positions(self, shape): + _, channels, frames, mel_bins = shape.audio_latent + return self.audio_patchifier.get_patch_grid_bounds( + output_shape=AudioLatentShape( + frames=frames, + mel_bins=mel_bins, + batch=shape.audio_latent[0], + channels=channels, + ), + device=self.device, + ) + + @staticmethod + def _singleton_shape(shape, name): + shape = tuple(int(dimension) for dimension in shape) + expected = 4 if name.startswith("audio") else 5 + without_batch = expected - 1 + if len(shape) == without_batch: + return (1, *shape) + if len(shape) == expected: + if shape[0] != 1: + raise ValueError( + f"training.dmd.{name} must begin with singleton dimension 1, got {shape}." + ) + return shape + raise ValueError(f"training.dmd.{name} has unsupported shape {shape}.") diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py new file mode 100644 index 000000000..3d3ff3a59 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py @@ -0,0 +1,204 @@ +"""Flow-matching capability for LTX joint video/audio models.""" + +from __future__ import annotations + +from typing import Any, Mapping + +import torch +from torch import Tensor + +from lightx2v_train.model_capabilities import ( + BoundCapability, + FlowMatchingSFTCapability, + LossResult, + SFTStepContext, +) +from lightx2v_train.model_zoo.native.ltx2 import ( + AudioLatentShape, + AudioPatchifier, + Modality, + SpatioTemporalScaleFactors, + VideoLatentPatchifier, + VideoLatentShape, + get_pixel_coords, +) + + +class LTXFlowMatchingCapability(BoundCapability, FlowMatchingSFTCapability): + """Supervised flow matching for LTX's joint video/audio denoiser.""" + + def __init__(self, model) -> None: + super().__init__(model) + config = model.config["training"].get("ltx2", {}) + self.video_loss_weight = float(config.get("video_loss_weight", 1.0)) + self.audio_loss_weight = float(config.get("audio_loss_weight", 1.0)) + self.default_fps = float(config.get("default_fps", 24.0)) + self.video_patchifier = VideoLatentPatchifier(patch_size=1) + self.audio_patchifier = AudioPatchifier(patch_size=1) + + def compute_loss( + self, + batch: Mapping[str, Any], + context: SFTStepContext, + ) -> LossResult: + inputs = batch["inputs"] + video_data = inputs.get("video_latents") + audio_data = inputs.get("audio_latents") + conditions = batch["conditioning"].get("positive") + if video_data is None or audio_data is None or conditions is None: + raise KeyError("LTX flow matching expects inputs.video_latents, inputs.audio_latents, and conditioning.positive.") + + video_latents = self._video_latents(video_data, context.running_dtype) + audio_latents = self._audio_latents(audio_data, context.running_dtype) + video_tokens = self.video_patchifier.patchify(video_latents) + audio_tokens = self.audio_patchifier.patchify(audio_latents) + + video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) + video_noisy = self._initialize_noisy_target( + video_tokens, + context.noise_scheduler, + ) + audio_noisy = self._initialize_noisy_target( + audio_tokens, + context.noise_scheduler, + sigmas=video_noisy[-1], + ) + video_input, video_target, video_timesteps, video_mask, video_sigmas = video_noisy + audio_input, audio_target, audio_timesteps, audio_mask, audio_sigmas = audio_noisy + + video_positions = self._video_positions( + video_data, + video_latents, + ) + audio_positions = self.audio_patchifier.get_patch_grid_bounds( + output_shape=AudioLatentShape( + frames=audio_tokens.shape[1], + mel_bins=16, + batch=1, + channels=8, + ), + device=audio_tokens.device, + ) + video = Modality( + enabled=True, + latent=video_input, + sigma=video_sigmas, + timesteps=video_timesteps, + positions=video_positions, + context=video_context, + context_mask=context_mask, + ) + audio = Modality( + enabled=True, + latent=audio_input, + sigma=audio_sigmas, + timesteps=audio_timesteps, + positions=audio_positions, + context=audio_context, + context_mask=context_mask, + ) + + with self.model.transformer_forward_context(): + video_prediction, audio_prediction = self.model.denoiser_module()( + video=video, + audio=audio, + perturbations=None, + ) + + video_loss = self._masked_mse( + video_prediction, + video_target, + video_mask, + ).mean() + audio_loss = self._masked_mse( + audio_prediction, + audio_target, + audio_mask, + ).mean() + loss = self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss + return LossResult( + loss=loss, + metrics={ + "video_loss": video_loss.detach(), + "audio_loss": audio_loss.detach(), + }, + ) + + def _video_latents(self, data, dtype): + latents = data["latents"].to(device=self.model.device, dtype=dtype) + if latents.ndim == 4: + latents = latents.unsqueeze(0) + if latents.ndim != 5: + raise ValueError(f"LTX video latents must have shape [B,C,F,H,W] or [C,F,H,W], got {tuple(latents.shape)}.") + if latents.shape[0] != 1: + raise ValueError("LTX flow matching only supports physical batch size 1.") + return latents.contiguous() + + def _audio_latents(self, data, dtype): + latents = data["latents"].to(device=self.model.device, dtype=dtype) + if latents.ndim == 3: + latents = latents.unsqueeze(0) + if latents.ndim != 4: + raise ValueError(f"LTX audio latents must have shape [B,C,T,F] or [C,T,F], got {tuple(latents.shape)}.") + if latents.shape[0] != 1: + raise ValueError("LTX flow matching only supports physical batch size 1.") + return latents.contiguous() + + @staticmethod + def _initialize_noisy_target(latents: Tensor, scheduler, sigmas=None): + _, sequence_length, _ = latents.shape + if sigmas is None: + sigmas = scheduler.sample_timestep_or_sigma( + seq_len=sequence_length, + ) + sigmas = sigmas.to(device=latents.device, dtype=latents.dtype) + noise = torch.randn_like(latents) + expanded = sigmas.view(-1, 1, 1) + noisy = (1 - expanded) * latents + expanded * noise + target = noise - latents + timesteps = ( + sigmas.view(-1, 1) + .expand( + 1, + sequence_length, + ) + .clone() + ) + mask = torch.ones( + 1, + sequence_length, + dtype=torch.bool, + device=latents.device, + ) + return noisy, target, timesteps, mask, sigmas + + def _video_positions(self, data, latents): + frames = int(data["num_frames"][0].item()) if "num_frames" in data else latents.shape[2] + height = int(data["height"][0].item()) if "height" in data else latents.shape[3] + width = int(data["width"][0].item()) if "width" in data else latents.shape[4] + fps = float(data["fps"][0].item()) if "fps" in data else self.default_fps + latent_coords = self.video_patchifier.get_patch_grid_bounds( + output_shape=VideoLatentShape( + frames=frames, + height=height, + width=width, + batch=1, + channels=128, + ), + device=latents.device, + ) + positions = get_pixel_coords( + latent_coords=latent_coords, + scale_factors=SpatioTemporalScaleFactors.default(), + causal_fix=True, + ).float() + positions[:, 0, ...] /= fps + return positions + + @staticmethod + def _masked_mse(prediction, target, mask): + target_length = target.shape[1] + prediction = prediction[:, -target_length:, :] + weights = mask[:, -target_length:].unsqueeze(-1).float() + loss = (prediction.float() - target.float()).square() * weights + return loss.mean(dim=(-2, -1)) / weights.mean(dim=(-2, -1)).clamp_min(1e-8) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py new file mode 100644 index 000000000..f1c7767b7 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py @@ -0,0 +1,357 @@ +"""Teacher-forcing capability for LTX models.""" + +from __future__ import annotations + +from typing import Any, Mapping + +import torch + +from lightx2v_train.model_capabilities import ( + LossResult, + TeacherForcingCapability, + TeacherForcingStepContext, +) +from lightx2v_train.model_zoo.native.ltx2 import AudioLatentShape, Modality +from lightx2v_train.runtime.distributed import is_sequence_parallel_enabled + +from .ltx_flow_matching_capability import LTXFlowMatchingCapability + + +class LTXTeacherForcingCapability( + LTXFlowMatchingCapability, + TeacherForcingCapability, +): + """Chunk-wise teacher forcing for LTX's joint video/audio denoiser.""" + + def __init__(self, model) -> None: + super().__init__(model) + self._mask_cache = {} + + def compute_loss( + self, + batch: Mapping[str, Any], + context: TeacherForcingStepContext, + ) -> LossResult: + if is_sequence_parallel_enabled(): + raise ValueError("LTX teacher forcing does not support sequence parallel.") + + inputs = batch["inputs"] + video_data = inputs.get("video_latents") + audio_data = inputs.get("audio_latents") + conditions = batch["conditioning"].get("positive") + if video_data is None or audio_data is None or conditions is None: + raise KeyError("LTX teacher forcing expects inputs.video_latents, inputs.audio_latents, and conditioning.positive.") + + with torch.no_grad(): + video_latents = self._video_latents( + video_data, + context.running_dtype, + ) + audio_latents = self._audio_latents( + audio_data, + context.running_dtype, + ) + video_latents, audio_latents = self._crop_to_chunk_multiple( + video_latents, + audio_latents, + context.num_frame_per_chunk, + ) + num_video_frames = int(video_latents.shape[2]) + video_tokens = self.video_patchifier.patchify(video_latents) + audio_tokens = self.audio_patchifier.patchify(audio_latents) + video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) + + video_inputs = self._prepare_video_inputs( + video_tokens, + num_frames=num_video_frames, + spatial_tokens=int(video_latents.shape[3] * video_latents.shape[4]), + context=context, + ) + ( + video_input, + video_target, + video_timesteps, + video_weights, + video_sigmas, + video_frame_weights, + video_augmentation_sigmas, + video_attention_mask, + video_block_ids, + ) = video_inputs + audio_inputs = self._prepare_audio_inputs( + audio_tokens, + frame_sigmas=video_sigmas, + frame_weights=video_frame_weights, + augmentation_frame_sigmas=video_augmentation_sigmas, + num_video_frames=num_video_frames, + frames_per_chunk=context.num_frame_per_chunk, + ) + ( + audio_input, + audio_target, + audio_timesteps, + audio_weights, + audio_sigmas, + audio_attention_mask, + audio_block_ids, + ) = audio_inputs + + video_positions = self._video_positions( + video_data, + video_latents, + ) + video_positions = torch.cat( + [video_positions, video_positions], + dim=2, + ) + audio_positions = self.audio_patchifier.get_patch_grid_bounds( + output_shape=AudioLatentShape( + frames=audio_tokens.shape[1], + mel_bins=16, + batch=1, + channels=8, + ), + device=audio_tokens.device, + ) + audio_positions = torch.cat( + [audio_positions, audio_positions], + dim=2, + ) + video_cross_mask = self._cross_attention_mask( + video_tokens.device, + video_block_ids, + audio_block_ids, + ) + audio_cross_mask = self._cross_attention_mask( + audio_tokens.device, + audio_block_ids, + video_block_ids, + ) + + video = Modality( + enabled=True, + latent=video_input, + sigma=video_sigmas[:, 0], + timesteps=video_timesteps, + positions=video_positions, + context=video_context, + context_mask=context_mask, + attention_mask=video_attention_mask, + cross_attention_mask=video_cross_mask, + ) + audio = Modality( + enabled=True, + latent=audio_input, + sigma=audio_sigmas[:, 0], + timesteps=audio_timesteps, + positions=audio_positions, + context=audio_context, + context_mask=context_mask, + attention_mask=audio_attention_mask, + cross_attention_mask=audio_cross_mask, + ) + with self.model.transformer_forward_context(): + video_prediction, audio_prediction = self.model.denoiser_module()( + video=video, + audio=audio, + perturbations=None, + ) + + video_loss = self._weighted_token_mse( + video_prediction, + video_target, + video_weights, + ).mean() + audio_loss = self._weighted_token_mse( + audio_prediction, + audio_target, + audio_weights, + ).mean() + loss = self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss + return LossResult( + loss=loss, + metrics={ + "video_loss": video_loss.detach(), + "audio_loss": audio_loss.detach(), + }, + ) + + @staticmethod + def _crop_to_chunk_multiple(video, audio, frames_per_chunk): + if frames_per_chunk <= 1: + return video, audio + num_frames = int(video.shape[2]) + keep_frames = num_frames // frames_per_chunk * frames_per_chunk + if keep_frames == num_frames: + return video, audio + if keep_frames <= 0: + raise ValueError(f"LTX sample has {num_frames} latent frames, fewer than frames_per_chunk={frames_per_chunk}.") + audio_frames = int(audio.shape[2]) + keep_audio_frames = max(1, audio_frames * keep_frames // num_frames) + return ( + video[:, :, :keep_frames].contiguous(), + audio[:, :, :keep_audio_frames].contiguous(), + ) + + def _prepare_video_inputs( + self, + tokens, + num_frames, + spatial_tokens, + context, + ): + frame_sigmas, frame_weights = context.scheduler.sample_chunkwise( + num_frames=num_frames, + num_frame_per_chunk=context.num_frame_per_chunk, + device=tokens.device, + dtype=tokens.dtype, + ) + token_sigmas = frame_sigmas.repeat_interleave(spatial_tokens, dim=1) + token_weights = frame_weights.repeat_interleave(spatial_tokens, dim=1) + noisy, target, noise = self._add_token_noise(tokens, token_sigmas) + clean = tokens + clean_timesteps = torch.zeros_like(token_sigmas) + augmentation_frame_sigmas = None + if context.noise_augmentation_max_timestep > 0: + augmentation_frame_sigmas = context.scheduler.sample_clean_augmentation( + num_frames=num_frames, + num_frame_per_chunk=context.num_frame_per_chunk, + max_timestep=context.noise_augmentation_max_timestep, + device=tokens.device, + dtype=tokens.dtype, + ) + clean_timesteps = augmentation_frame_sigmas.repeat_interleave( + spatial_tokens, + dim=1, + ) + clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise + frame_ids = torch.arange( + num_frames, + device=tokens.device, + ).repeat_interleave(spatial_tokens) + block_ids = frame_ids // context.num_frame_per_chunk + attention_mask = self._attention_mask(tokens.device, block_ids) + return ( + torch.cat([clean, noisy], dim=1), + target, + torch.cat([clean_timesteps, token_sigmas], dim=1), + token_weights, + frame_sigmas, + frame_weights, + augmentation_frame_sigmas, + attention_mask, + block_ids, + ) + + def _prepare_audio_inputs( + self, + tokens, + frame_sigmas, + frame_weights, + augmentation_frame_sigmas, + num_video_frames, + frames_per_chunk, + ): + _, sequence_length, _ = tokens.shape + frame_ids = (torch.arange(sequence_length, device=tokens.device) * max(1, num_video_frames) // max(1, sequence_length)).clamp(max=frame_sigmas.shape[1] - 1) + gather_ids = frame_ids.unsqueeze(0) + token_sigmas = frame_sigmas.gather(1, gather_ids) + token_weights = frame_weights.gather(1, gather_ids) + noisy, target, noise = self._add_token_noise(tokens, token_sigmas) + clean = tokens + clean_timesteps = torch.zeros_like(token_sigmas) + if augmentation_frame_sigmas is not None: + clean_timesteps = augmentation_frame_sigmas.gather(1, gather_ids) + clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise + block_ids = frame_ids // frames_per_chunk + attention_mask = self._attention_mask(tokens.device, block_ids) + return ( + torch.cat([clean, noisy], dim=1), + target, + torch.cat([clean_timesteps, token_sigmas], dim=1), + token_weights, + token_sigmas, + attention_mask, + block_ids, + ) + + def _attention_mask(self, device, block_ids): + key = ("self", self._device_key(device), self._block_key(block_ids)) + if key not in self._mask_cache: + self._mask_cache[key] = self._call_mask_builder( + "prepare_teacher_forcing_mask_from_block_ids", + device, + block_ids, + ) + return self._mask_cache[key] + + def _cross_attention_mask(self, device, query_ids, key_ids): + key = ( + "cross", + self._device_key(device), + self._block_key(query_ids), + self._block_key(key_ids), + ) + if key not in self._mask_cache: + self._mask_cache[key] = self._call_mask_builder( + "prepare_teacher_forcing_cross_mask_from_block_ids", + device, + query_ids, + key_ids, + ) + return self._mask_cache[key] + + def _call_mask_builder(self, method_name, *args): + transformer = self.model.denoiser_module() + candidates = ( + transformer, + getattr(transformer, "module", None), + getattr(transformer, "_fsdp_wrapped_module", None), + ) + for candidate in candidates: + method = getattr(candidate, method_name, None) + if method is not None: + return method(*args) + raise RuntimeError(f"LTX transformer does not implement {method_name}().") + + @staticmethod + def _device_key(device): + device = torch.device(device) + index = device.index + if device.type == "cuda" and index is None and torch.cuda.is_available(): + index = torch.cuda.current_device() + return device.type, index + + @staticmethod + def _block_key(block_ids): + values = block_ids.detach().to(dtype=torch.long) + if values.numel() == 0: + return 0, None, None, None, None, None + positions = torch.arange( + values.numel(), + device=values.device, + dtype=torch.long, + ) + return ( + int(values.numel()), + int(values[0].item()), + int(values[-1].item()), + int(values.max().item()), + int(values.sum().item()), + int((values * positions).sum().item()), + ) + + @staticmethod + def _add_token_noise(tokens, sigmas, noise=None): + if noise is None: + noise = torch.randn_like(tokens) + expanded = sigmas.unsqueeze(-1) + noisy = (1.0 - expanded) * tokens + expanded * noise + return noisy, noise - tokens, noise + + @staticmethod + def _weighted_token_mse(prediction, target, weights): + prediction = prediction[:, -target.shape[1] :, :] + loss = (prediction.float() - target.float()).square().mean(dim=-1) + weights = weights.float() + return (loss * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1e-8) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py similarity index 92% rename from lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py rename to lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py index 1fcc3061e..57b59c7d7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx_t2av.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py @@ -9,6 +9,18 @@ from safetensors import safe_open from safetensors.torch import load_file +from lightx2v_train.model_capabilities import ( + AutoregressiveDistillationCapability, + DistillationCapability, + FlowMatchingSFTCapability, + TeacherForcingCapability, +) +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_autoregressive_distillation_capability import ( + LTXAutoregressiveDistillationCapability, +) +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_distillation_capability import LTXDistillationCapability +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_flow_matching_capability import LTXFlowMatchingCapability +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_teacher_forcing_capability import LTXTeacherForcingCapability from lightx2v_train.model_zoo.native.ltx2 import ( EMBEDDINGS_PROCESSOR_KEY_OPS, GEMMA_LLM_KEY_OPS, @@ -25,7 +37,7 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype -from .base import BaseModel +from ..base import BaseModel @MODEL_REGISTER("ltx_t2av_ar") @@ -33,6 +45,26 @@ class LTX2T2AVModel(BaseModel): pipeline_cls = None + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + LTXFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + LTXDistillationCapability(self), + ) + if self.config["model"]["name"] == "ltx_t2av_ar": + self.capabilities.register( + AutoregressiveDistillationCapability, + LTXAutoregressiveDistillationCapability(self), + ) + self.capabilities.register( + TeacherForcingCapability, + LTXTeacherForcingCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): model_config = self.config["model"] model_path = model_config["pretrained_model_name_or_path"] @@ -346,6 +378,8 @@ def encode_prompt_condition(self, prompt): raise RuntimeError("LTX2 embeddings processor feature_extractor is required for prompt-based training.") prompts = [prompt] if isinstance(prompt, str) else list(prompt) + if len(prompts) != 1: + raise ValueError(f"LTX2 training requires exactly one prompt, got {len(prompts)}.") with torch.no_grad(): encoded = self.text_encoder.encode(prompts, padding_side="left") diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py index d05773b44..b4b317564 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py @@ -320,13 +320,13 @@ def pre_dit( raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") batch_size = action_tokens.shape[0] + if self.training and batch_size != 1: + raise ValueError("FastWAM action training only supports physical batch size 1.") if context.shape[0] != batch_size: raise ValueError(f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}") if timestep.shape[0] not in (1, batch_size): raise ValueError(f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}") if timestep.shape[0] == 1 and batch_size > 1: - if self.training: - raise ValueError("During training, action timestep length must match batch_size.") timestep = timestep.expand(batch_size) if context_mask is None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py index 804eef573..023665eb3 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py @@ -296,7 +296,9 @@ def build_inputs(self, sample, tiled: bool = False): if video.shape[1] != 3: raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") - batch_size, _, num_frames, height, width = video.shape + _, _, num_frames, height, width = video.shape + if video.shape[0] != 1: + raise ValueError("FastWAM training only supports physical batch size 1.") if height % 16 != 0 or width % 16 != 0: raise ValueError(f"Video spatial dims must be multiples of 16, got H={height}, W={width}") if num_frames % 4 != 1: @@ -318,15 +320,15 @@ def build_inputs(self, sample, tiled: bool = False): if action_is_pad is not None: if action_is_pad.ndim != 2: raise ValueError(f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}") - if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: - raise ValueError(f"`sample['action_is_pad']` shape mismatch: got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})") + if action_is_pad.shape[0] != 1 or action_is_pad.shape[1] != action_horizon: + raise ValueError(f"`sample['action_is_pad']` shape mismatch: got {tuple(action_is_pad.shape)} vs expected (1, {action_horizon})") image_is_pad = sample.get("image_is_pad", None) if image_is_pad is not None: if image_is_pad.ndim != 2: raise ValueError(f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}") - if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: - raise ValueError(f"`sample['image_is_pad']` shape mismatch: got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})") + if image_is_pad.shape[0] != 1 or image_is_pad.shape[1] != num_frames: + raise ValueError(f"`sample['image_is_pad']` shape mismatch: got {tuple(image_is_pad.shape)} vs expected (1, {num_frames})") input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) input_latents = self._encode_video_latents(input_video, tiled=tiled) @@ -432,7 +434,8 @@ def _compute_video_loss_per_sample( def training_loss(self, sample, tiled: bool = False): inputs = self.build_inputs(sample, tiled=tiled) input_latents = inputs["input_latents"] - batch_size = input_latents.shape[0] + if input_latents.shape[0] != 1: + raise ValueError("FastWAM training only supports physical batch size 1.") context = inputs["context"] context_mask = inputs["context_mask"] action = inputs["action"] @@ -441,7 +444,6 @@ def training_loss(self, sample, tiled: bool = False): noise_video = torch.randn_like(input_latents) timestep_video = self.train_video_scheduler.sample_training_t( - batch_size=batch_size, device=self.device, dtype=input_latents.dtype, ) @@ -453,7 +455,6 @@ def training_loss(self, sample, tiled: bool = False): noise_action = torch.randn_like(action) timestep_action = self.train_action_scheduler.sample_training_t( - batch_size=batch_size, device=self.device, dtype=action.dtype, ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py index 97f3aadd7..ac756f593 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py @@ -80,6 +80,8 @@ def _validate_inputs(self, x, timestep, context, context_mask): if timestep.ndim != 1: raise ValueError(f"timestep must be [B] or [1], got {tuple(timestep.shape)}") batch_size = x.shape[0] + if self.training and batch_size != 1: + raise ValueError("FastWAM video training only supports physical batch size 1.") if context.shape[0] != batch_size: if not self.training and batch_size == 1: x = x.expand(context.shape[0], -1, -1, -1, -1) @@ -89,8 +91,6 @@ def _validate_inputs(self, x, timestep, context, context_mask): if timestep.shape[0] not in (1, batch_size): raise ValueError(f"timestep length must be 1 or {batch_size}, got {timestep.shape[0]}") if timestep.shape[0] == 1 and batch_size > 1: - if self.training: - raise ValueError("Training timestep length must match batch size.") timestep = timestep.expand(batch_size) if context_mask is None: context_mask = torch.ones(context.shape[:2], dtype=torch.bool, device=context.device) diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/__init__.py new file mode 100644 index 000000000..5da1a1199 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/__init__.py @@ -0,0 +1 @@ +"""Qwen Image model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py similarity index 93% rename from lightx2v_train/lightx2v_train/model_zoo/qwen_image.py rename to lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index d3fea8380..c13c0443d 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -6,9 +6,11 @@ from diffusers.image_processor import VaeImageProcessor from torch import nn +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericConsistencyCapability, GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER -from .base import BaseModel +from ..base import BaseModel @dataclass @@ -28,6 +30,21 @@ class QwenImageModel(BaseModel): pipeline_cls = QwenImagePipeline + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + GenericDistillationCapability(self), + ) + self.capabilities.register( + ConsistencyCapability, + GenericConsistencyCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): if transformer_only: if reference_model is not None: @@ -197,9 +214,9 @@ def prepare_infer_latents(self, height, width, generator=None): shape = (1, self.vae.config.z_dim, 1, 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): + def dmd_latent_shape(self, height, width): return ( - int(batch_size), + 1, int(self.vae.config.z_dim), 1, int(height) // self.vae_scale_factor, diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py similarity index 96% rename from lightx2v_train/lightx2v_train/model_zoo/qwen_image_edit.py rename to lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index bd2dc887d..00388f80b 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -6,9 +6,11 @@ from diffusers import AutoencoderKLQwenImage, QwenImageEditPlusPipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor +from lightx2v_train.model_capabilities import FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER -from .base import BaseModel +from ..base import BaseModel CONDITION_IMAGE_AREA = 384 * 384 @@ -42,6 +44,13 @@ class QwenImageEditModel(BaseModel): pipeline_cls = QwenImageEditPlusPipeline + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): if transformer_only: if reference_model is not None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/wan/__init__.py new file mode 100644 index 000000000..d479e5ae9 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/__init__.py @@ -0,0 +1 @@ +"""Wan-family model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py new file mode 100644 index 000000000..bc9c89a4e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py @@ -0,0 +1 @@ +"""Capabilities implemented specifically by Wan-family models.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py new file mode 100644 index 000000000..75de6798b --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py @@ -0,0 +1,303 @@ +"""Autoregressive distillation capability for causal Wan models.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from lightx2v_train.model_capabilities import ( + AutoregressiveDistillationCapability, + AutoregressiveRolloutContext, + BoundCapability, +) +from lightx2v_train.runtime.distributed import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + is_sequence_parallel_enabled, +) +from lightx2v_train.runtime.sequence_parallel import all_gather_sequence + + +class WanAutoregressiveDistillationCapability( + BoundCapability, + AutoregressiveDistillationCapability, +): + """Cached chunk-wise rollout for causal Wan transformers.""" + + def rollout( + self, + condition, + latent_shape, + initial_latents, + context: AutoregressiveRolloutContext, + ): + latents = initial_latents + _, _, num_frames, _, _ = latents.shape + if latents.shape[0] != 1: + raise ValueError("Wan autoregressive DMD only supports physical batch size 1.") + if num_frames % context.frames_per_chunk: + raise ValueError(f"Autoregressive latent frames ({num_frames}) must be divisible by frames_per_chunk ({context.frames_per_chunk}).") + + use_sp_cache = bool(context.sequence_parallel_cache and is_sequence_parallel_enabled()) + parallel_size = get_sequence_parallel_world_size() if use_sp_cache else 1 + parallel_rank = get_sequence_parallel_rank() if use_sp_cache else 0 + if use_sp_cache and context.frames_per_chunk % parallel_size: + raise ValueError("frames_per_chunk must be divisible by sequence parallel size when sp_cache is enabled.") + transformer = self.model.denoiser_module() + if use_sp_cache and bool(getattr(transformer, "defer_kv_cache_updates", False)): + raise ValueError("sp_cache is incompatible with deferred KV-cache updates.") + + transformer.train() + model_context = self.model._condition_to_context_tensor( + condition, + batch_size=1, + ) + frame_sequence_length = self._frame_sequence_length(latents) + kv_cache, cross_attention_cache = self._new_caches( + latents.dtype, + latents.device, + num_frames, + frame_sequence_length, + sequence_parallel_cache=use_sp_cache, + ) + num_blocks = num_frames // context.frames_per_chunk + exit_indices = self._sample_exit_indices( + 1 if context.same_step_across_blocks else num_blocks, + len(context.denoising_steps), + latents.device, + ) + + output_chunks = [] + current_start_frame = 0 + for block_index in range(num_blocks): + if use_sp_cache: + local_frames = context.frames_per_chunk // parallel_size + local_start = current_start_frame + parallel_rank * local_frames + else: + local_frames = context.frames_per_chunk + local_start = current_start_frame + block_latents = latents[:, :, local_start : local_start + local_frames] + exit_index = int(exit_indices[0] if context.same_step_across_blocks else exit_indices[block_index]) + + x0 = None + for step_index, current_timestep in enumerate(context.denoising_steps): + timestep = torch.full( + (1, local_frames), + float(current_timestep), + device=latents.device, + dtype=torch.float32, + ) + enable_grad = context.grad_enabled and step_index == exit_index + grad_context = torch.enable_grad if enable_grad else torch.no_grad + with grad_context(): + flow = self._forward_chunk( + block_latents, + timestep, + model_context, + kv_cache, + cross_attention_cache, + current_start=current_start_frame * frame_sequence_length, + cache_start=current_start_frame * frame_sequence_length, + local_frame_offset=local_start, + balanced_sequence_parallel=use_sp_cache, + ) + x0 = self._flow_to_x0( + block_latents, + flow, + timestep, + context.denoising_scheduler, + ) + if step_index == exit_index: + break + next_timestep = torch.full( + (1, local_frames), + float(context.denoising_steps[step_index + 1]), + device=latents.device, + dtype=torch.float32, + ) + with torch.no_grad(): + block_latents = self._add_noise( + x0, + torch.randn_like(x0), + next_timestep, + context.denoising_scheduler, + ) + + output_chunks.append(all_gather_sequence(x0, dim=2) if use_sp_cache else x0) + cache_latents = x0.detach() + cache_timestep = torch.full( + (1, local_frames), + context.context_noise, + device=latents.device, + dtype=torch.float32, + ) + if context.context_noise > 0: + cache_latents = self._add_noise( + cache_latents, + torch.randn_like(cache_latents), + cache_timestep, + context.denoising_scheduler, + ) + with torch.no_grad(): + self._forward_chunk( + cache_latents, + cache_timestep, + model_context, + kv_cache, + cross_attention_cache, + current_start=current_start_frame * frame_sequence_length, + cache_start=current_start_frame * frame_sequence_length, + local_frame_offset=local_start, + balanced_sequence_parallel=use_sp_cache, + ) + current_start_frame += context.frames_per_chunk + + generated = torch.cat(output_chunks, dim=2).to(dtype=context.running_dtype) + return generated, exit_indices + + def _forward_chunk( + self, + latents, + timestep, + condition, + kv_cache, + cross_attention_cache, + **kwargs, + ): + with self.model.transformer_forward_context(): + return self.model.denoiser_module()( + latents, + t=timestep, + context=condition, + seq_len=self.model._sequence_length(latents), + kv_cache=kv_cache, + crossattn_cache=cross_attention_cache, + **kwargs, + ) + + def _new_caches( + self, + dtype, + device, + num_frames, + frame_sequence_length, + sequence_parallel_cache, + ): + transformer = self.model.denoiser_module() + num_layers = int(getattr(transformer, "num_layers", len(transformer.blocks))) + num_heads = int(transformer.num_heads) + head_dim = int(transformer.dim // transformer.num_heads) + kv_heads = num_heads + if sequence_parallel_cache: + parallel_size = get_sequence_parallel_world_size() + if num_heads % parallel_size: + raise ValueError(f"Transformer heads ({num_heads}) must be divisible by sequence parallel size ({parallel_size}).") + kv_heads = num_heads // parallel_size + local_attention_size = int(getattr(transformer, "local_attn_size", -1)) + cache_size = num_frames * frame_sequence_length if local_attention_size == -1 else local_attention_size * frame_sequence_length + kv_cache = [ + { + "k": torch.zeros( + (1, cache_size, kv_heads, head_dim), + dtype=dtype, + device=device, + ), + "v": torch.zeros( + (1, cache_size, kv_heads, head_dim), + dtype=dtype, + device=device, + ), + "global_end_index": torch.zeros( + 1, + dtype=torch.long, + device=device, + ), + "local_end_index": torch.zeros( + 1, + dtype=torch.long, + device=device, + ), + } + for _ in range(num_layers) + ] + cross_attention_cache = [ + { + "k": torch.zeros( + ( + 1, + self.model.max_sequence_length, + num_heads, + head_dim, + ), + dtype=dtype, + device=device, + ), + "v": torch.zeros( + ( + 1, + self.model.max_sequence_length, + num_heads, + head_dim, + ), + dtype=dtype, + device=device, + ), + "is_init": False, + } + for _ in range(num_layers) + ] + return kv_cache, cross_attention_cache + + def _frame_sequence_length(self, latent): + _, _, _, height, width = latent.shape + patch_t, patch_h, patch_w = self.model.patch_size + if patch_t != 1: + raise ValueError(f"Autoregressive rollout requires temporal patch size 1, got {patch_t}.") + return height * width // (patch_h * patch_w) + + @staticmethod + def _sample_exit_indices(count, num_steps, device): + if dist.is_available() and dist.is_initialized(): + if dist.get_rank() == 0: + indices = torch.randint(0, num_steps, (count,), device=device) + else: + indices = torch.empty(count, dtype=torch.long, device=device) + dist.broadcast(indices, src=0) + return indices.tolist() + return torch.randint(0, num_steps, (count,), device=device).tolist() + + @staticmethod + def _sigma_from_timestep(timestep, dtype, scheduler): + timesteps = scheduler.timesteps.to( + device=timestep.device, + dtype=torch.float32, + ) + sigmas = scheduler.sigmas.to(device=timestep.device, dtype=dtype) + flat = timestep.flatten().float() + indices = torch.argmin( + (timesteps.unsqueeze(0) - flat.unsqueeze(1)).abs(), + dim=1, + ) + return sigmas[indices].reshape(timestep.shape) + + @classmethod + def _flow_to_x0(cls, sample, flow, timestep, scheduler): + sigma = cls._sigma_from_timestep(timestep, sample.dtype, scheduler) + sigma = sigma.reshape( + sigma.shape[0], + 1, + sigma.shape[1], + *([1] * (sample.ndim - 3)), + ) + return (sample - sigma * flow).to(dtype=sample.dtype) + + @classmethod + def _add_noise(cls, x0, noise, timestep, scheduler): + sigma = cls._sigma_from_timestep(timestep, x0.dtype, scheduler) + sigma = sigma.reshape( + sigma.shape[0], + 1, + sigma.shape[1], + *([1] * (x0.ndim - 3)), + ) + return ((1.0 - sigma) * x0 + sigma * noise).to(dtype=x0.dtype) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py new file mode 100644 index 000000000..466628602 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py @@ -0,0 +1,123 @@ +"""Distribution-matching capabilities for Wan-family video models.""" + +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistillationCapability, + _negative_prompt, + _require_single_prompt, + _require_singleton_tensor, +) +from lightx2v_train.utils.constants import ( + LINGBOT_VIDEO_NEGATIVE_PROMPT, + WAN_NEGATIVE_PROMPT, +) + + +class WanDistillationCapability(GenericDistillationCapability): + """Distribution-matching operations for Wan-family video models.""" + + @property + def default_negative_prompt(self): + return WAN_NEGATIVE_PROMPT + + @property + def default_lora_target_modules(self): + return ("q", "k", "v", "o", "ffn.0", "ffn.2") + + def latent_shape( + self, + batch, + shape_config, + image_sizes, + broadcast, + ): + del image_sizes, broadcast + prompt = batch["conditioning"].get("prompt", "") + _require_single_prompt(prompt) + configured = shape_config.get("image_or_video_shape") + if configured is not None: + shape = tuple(int(dimension) for dimension in configured) + if len(shape) != 5 or shape[0] != 1: + raise ValueError( + "training.dmd.image_or_video_shape must be a five-dimensional " + f"singleton shape beginning with 1, got {shape}." + ) + return shape + missing = [key for key in ("height", "width", "num_frames") if key not in shape_config] + if missing: + raise ValueError("Video DMD shape is missing: " + ", ".join(missing) + ".") + latent_frames = (int(shape_config["num_frames"]) - 1) // self.model.vae_scale_factor_temporal + 1 + return ( + 1, + self.model._latent_channels(), + latent_frames, + int(shape_config["height"]) // self.model.vae_scale_factor_spatial, + int(shape_config["width"]) // self.model.vae_scale_factor_spatial, + ) + + def encode_conditions( + self, + batch, + negative_prompt, + guidance_scale, + broadcast, + ): + conditioning = batch["conditioning"] + positive = conditioning.get("positive") + if positive is None: + return super().encode_conditions( + batch, + negative_prompt, + guidance_scale, + broadcast, + ) + with torch.no_grad(): + condition = self.prepare_cached_condition(positive) + if guidance_scale > 1: + cached_negative = conditioning.get("negative") + if cached_negative is not None: + negative = self.prepare_cached_condition(cached_negative) + else: + scalar = _require_single_prompt(conditioning.get("prompt", "")) + prompts = _negative_prompt( + conditioning, + negative_prompt, + scalar=scalar, + ) + negative = self.model.encode_prompt_condition(prompts) + else: + negative = None + return ( + broadcast(condition), + broadcast(negative) if negative is not None else None, + ) + + def prepare_cached_condition(self, condition): + if torch.is_tensor(condition): + prompt_embed = condition + elif isinstance(condition, dict) and "prompt_embed" in condition: + prompt_embed = condition["prompt_embed"] + else: + raise KeyError("Wan DMD cached condition expects a prompt_embed tensor.") + prompt_embed = prompt_embed.to( + device=self.device, + dtype=self.model.running_dtype, + ) + if prompt_embed.ndim == 2: + prompt_embed = prompt_embed.unsqueeze(0) + _require_singleton_tensor(prompt_embed, "Wan cached prompt embedding") + return {"prompt_embed": prompt_embed} + + +class LingBotDistillationCapability(WanDistillationCapability): + @property + def default_negative_prompt(self): + return LINGBOT_VIDEO_NEGATIVE_PROMPT + + @property + def default_lora_target_modules(self): + return None + + def prepare_cached_condition(self, condition): + return self.model.prepare_text_condition(condition) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_teacher_forcing_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_teacher_forcing_capability.py new file mode 100644 index 000000000..621fb28d8 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_teacher_forcing_capability.py @@ -0,0 +1,122 @@ +"""Teacher-forcing capability for causal Wan models.""" + +from typing import Any, Mapping + +import torch +import torch.distributed as dist + +from lightx2v_train.model_capabilities import ( + BoundCapability, + LossResult, + TeacherForcingCapability, + TeacherForcingStepContext, +) +from lightx2v_train.runtime.distributed import ( + get_sequence_parallel_group, + get_sequence_parallel_world_size, + is_sequence_parallel_enabled, +) +from lightx2v_train.runtime.sequence_parallel import ( + sequence_parallel_frame_slice, +) + + +class WanTeacherForcingCapability(BoundCapability, TeacherForcingCapability): + """Chunk-wise teacher forcing for a causal Wan denoiser.""" + + def compute_loss( + self, + batch: Mapping[str, Any], + context: TeacherForcingStepContext, + ) -> LossResult: + broadcast = context.broadcast + + with torch.no_grad(): + latent = self._latent(batch, context.running_dtype) + _, _, num_frames, _, _ = latent.shape + parallel_size = get_sequence_parallel_world_size() if is_sequence_parallel_enabled() else 1 + frame_multiple = parallel_size * context.num_frame_per_chunk + if num_frames % frame_multiple: + raise ValueError( + f"Teacher-forcing frames ({num_frames}) must be divisible by sequence_parallel_size * frames_per_chunk ({parallel_size} * {context.num_frame_per_chunk} = {frame_multiple})." + ) + + latent = broadcast(latent) + noise = broadcast(torch.randn_like(latent, dtype=context.running_dtype)) + sigmas, weights = context.scheduler.sample_chunkwise( + num_frames=num_frames, + num_frame_per_chunk=context.num_frame_per_chunk, + device=latent.device, + dtype=context.running_dtype, + ) + sigmas = broadcast(sigmas) + weights = broadcast(weights) + noisy_latent = context.scheduler.add_noise(latent, noise, sigmas) + condition = broadcast(self.model.encode_condition(batch)) + + clean_latent = latent + augmentation_sigmas = None + if context.noise_augmentation_max_timestep > 0: + augmentation_sigmas = context.scheduler.sample_clean_augmentation( + num_frames=num_frames, + num_frame_per_chunk=context.num_frame_per_chunk, + max_timestep=(context.noise_augmentation_max_timestep), + device=latent.device, + dtype=context.running_dtype, + ) + augmentation_sigmas = broadcast(augmentation_sigmas) + clean_latent = context.scheduler.add_noise( + latent, + noise, + augmentation_sigmas, + ) + + frame_start, frame_end, _ = sequence_parallel_frame_slice( + num_frames, + context.num_frame_per_chunk, + ) + latent = latent[:, :, frame_start:frame_end].contiguous() + noise = noise[:, :, frame_start:frame_end].contiguous() + noisy_latent = noisy_latent[:, :, frame_start:frame_end].contiguous() + clean_latent = clean_latent[:, :, frame_start:frame_end].contiguous() + sigmas = sigmas[:, frame_start:frame_end].contiguous() + weights = weights[:, frame_start:frame_end].contiguous() + if augmentation_sigmas is not None: + augmentation_sigmas = augmentation_sigmas[:, frame_start:frame_end].contiguous() + + prediction = self.model.denoise_teacher_forcing( + noisy_latent, + sigmas, + condition, + clean_latent=clean_latent, + aug_timestep_or_sigma=augmentation_sigmas, + frame_offset=frame_start, + global_num_frames=num_frames, + ) + target = noise - latent + frame_loss = (prediction.float() - target.float()).square().mean(dim=(1, 3, 4)) + weighted_loss = (frame_loss * weights).sum() + if is_sequence_parallel_enabled(): + dist.all_reduce( + weighted_loss, + op=dist.ReduceOp.SUM, + group=get_sequence_parallel_group(), + ) + return LossResult(loss=weighted_loss / num_frames) + + def _latent(self, batch, dtype): + latent = batch["inputs"].get("latents") + if latent is None: + latent = self.model.encode_to_latent(batch) + latent = latent.to(device=self.model.device, dtype=dtype) + if latent.ndim == 4: + latent = latent.unsqueeze(0) + if latent.shape[0] != 1: + raise ValueError("Wan teacher forcing only supports physical batch size 1.") + + channels = self.model._latent_channels() + if latent.shape[1] != channels and latent.shape[2] == channels: + latent = latent.permute(0, 2, 1, 3, 4).contiguous() + if latent.shape[1] != channels: + raise ValueError(f"Teacher-forcing latent channels ({latent.shape[1]}) do not match model channels ({channels}).") + return latent diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py new file mode 100644 index 000000000..b1d527b9e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py @@ -0,0 +1,40 @@ +"""World-action training capability for Wan FastWAM.""" + +from lightx2v_train.model_capabilities import ( + BoundCapability, + LossResult, + WorldActionTrainingCapability, +) + + +class WanWorldActionCapability( + BoundCapability, + WorldActionTrainingCapability, +): + def configure(self) -> None: + self.model.set_dit_only_trainable() + self.model.log_model_structure() + + def parameters(self): + return self.model.trainable_parameters() + + def module(self): + return self.model.unwrap_module() + + def compute_loss(self, batch, module=None): + if module is None: + module = self.module() + with self.model.autocast_context(): + loss, metrics = module(batch) + return LossResult(loss=loss, metrics=metrics) + + def evaluation_loss(self, batch): + with self.model.autocast_context(): + loss, metrics = self.module().training_loss(batch) + return LossResult(loss=loss, metrics=metrics) + + def load_checkpoint(self, path) -> None: + self.model.load_checkpoint(path) + + def save_checkpoint(self, path, step=None) -> None: + self.model.save_checkpoint(path, step=step) diff --git a/lightx2v_train/lightx2v_train/model_zoo/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py similarity index 92% rename from lightx2v_train/lightx2v_train/model_zoo/lingbot_video.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index 6411e7a87..174d22b54 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -9,10 +9,15 @@ from peft.utils import set_peft_model_state_dict from safetensors.torch import load_file +from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability +from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( + LingBotDistillationCapability, +) from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype -from .base import BaseModel +from ..base import BaseModel TOKEN_LENGTH = 37698 HIDDEN_STATE_SKIP_LAYER = 0 @@ -39,6 +44,18 @@ class LingBotVideoDenoiserInput: @MODEL_REGISTER("lingbot_video") class LingBotVideoModel(BaseModel): pipeline_cls = None + + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + LingBotDistillationCapability(self), + ) + vae_scale_factor_temporal = 4 vae_scale_factor_spatial = 8 @@ -77,7 +94,7 @@ def load_components(self, transformer_only=False, reference_model=None): def _load_transformer(self): try: - from .native.lingbot_video import LingBotVideoTransformer3DModel + from ..native.lingbot_video import LingBotVideoTransformer3DModel except ImportError as exc: raise ImportError("LingBot-Video training requires the diffusers APIs used by LingBotVideoTransformer3DModel.") from exc @@ -251,6 +268,8 @@ def encode_prompt_condition(self, prompt): if self.text_encoder is None or self.processor is None: raise RuntimeError("LingBot-Video text encoder is not loaded. Use cached conditions or set model.load_text_encoder=true.") prompts = [prompt] if isinstance(prompt, (str, dict)) else list(prompt) + if len(prompts) != 1: + raise ValueError(f"LingBot-Video training requires exactly one prompt, got {len(prompts)}.") prompts = [self._normalize_prompt(text) for text in prompts] texts = [self.prompt_template.format(text) for text in prompts] inputs = self.processor( @@ -320,6 +339,8 @@ def prepare_text_condition(self, condition): prompt_embed = prompt_embed.unsqueeze(0) if prompt_mask.ndim == 1: prompt_mask = prompt_mask.unsqueeze(0) + if prompt_embed.shape[0] != 1 or prompt_mask.shape[0] != 1: + raise ValueError("LingBot-Video cached conditions must have leading dimension 1.") return {"prompt_embed": prompt_embed, "prompt_attention_mask": prompt_mask} def encode_condition(self, sample): @@ -336,11 +357,13 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): hidden_states = denoiser_input.hidden_states.to(device=self.device, dtype=self.running_dtype) if hidden_states.ndim == 4: hidden_states = hidden_states.unsqueeze(0) + if hidden_states.shape[0] != 1: + raise ValueError("LingBot-Video training only supports physical batch size 1.") sigma = timestep_or_sigma.to(device=self.device, dtype=torch.float32) if sigma.ndim == 0: - sigma = sigma.expand(hidden_states.shape[0]) - elif sigma.shape[0] == 1 and hidden_states.shape[0] > 1: - sigma = sigma.expand(hidden_states.shape[0]) + sigma = sigma.reshape(1) + elif sigma.numel() != 1: + raise ValueError("LingBot-Video training requires exactly one sigma value.") timestep = sigma * 1000.0 prompt_embed = condition["prompt_embed"].to(device=self.device, dtype=self.running_dtype) prompt_mask = condition["prompt_attention_mask"].to(device=self.device) @@ -362,14 +385,14 @@ def cfg_on_denoiser_output(self): def _latent_channels(self): return int(self.transformer.config.in_channels) - def dmd_latent_shape(self, batch_size, height, width): + def dmd_latent_shape(self, height, width): num_frames = int(self.config.get("inference", {}).get("num_frames", 9)) if num_frames != 1 and (num_frames - 1) % self.vae_scale_factor_temporal != 0: raise ValueError(f"LingBot-Video num_frames must be 1 or 4n+1, got {num_frames}.") if int(height) % 16 != 0 or int(width) % 16 != 0: raise ValueError(f"LingBot-Video height and width must be multiples of 16, got {height}x{width}.") return ( - int(batch_size), + 1, self._latent_channels(), (num_frames - 1) // self.vae_scale_factor_temporal + 1, int(height) // self.vae_scale_factor_spatial, @@ -377,7 +400,7 @@ def dmd_latent_shape(self, batch_size, height, width): ) def prepare_infer_latents(self, height, width, generator=None): - shape = self.dmd_latent_shape(1, height, width) + shape = self.dmd_latent_shape(height, width) return torch.randn(shape, generator=generator, device=self.device, dtype=torch.float32) def _vae_latent_to_dit(self, latents): diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py similarity index 91% rename from lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py index 6b667bef7..517fe5d81 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py @@ -4,7 +4,14 @@ import torch from loguru import logger +from lightx2v_train.model_capabilities import ( + CapabilityProvider, + WorldActionTrainingCapability, +) from lightx2v_train.model_zoo.native.wan.fastwam import FastWAM +from lightx2v_train.model_zoo.wan.capability_adapters.wan_world_action_capability import ( + WanWorldActionCapability, +) from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype @@ -21,14 +28,21 @@ def _resolve_local_path(path, name, directory=False): @MODEL_REGISTER("wan_fastwam") -class WanFastWAMModel: +class WanFastWAMModel(CapabilityProvider): def __init__(self, config): + super().__init__() self.config = config self.model_config = config["model"] self.running_dtype = get_running_dtype(self.model_config.get("running_dtype", "bf16")) self.device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") self.module = None + def register_capabilities(self): + self.capabilities.register( + WorldActionTrainingCapability, + WanWorldActionCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): if transformer_only: raise ValueError("wan_fastwam does not support transformer_only loading.") diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_t2v.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py similarity index 90% rename from lightx2v_train/lightx2v_train/model_zoo/wan_t2v.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py index bd43cabc0..a0f328579 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan_t2v.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py @@ -11,15 +11,31 @@ from peft.utils import set_peft_model_state_dict from safetensors.torch import load_file +from lightx2v_train.model_capabilities import ( + AutoregressiveDistillationCapability, + DistillationCapability, + FlowMatchingSFTCapability, + TeacherForcingCapability, +) +from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability +from lightx2v_train.model_zoo.wan.capability_adapters.wan_autoregressive_distillation_capability import ( + WanAutoregressiveDistillationCapability, +) +from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( + WanDistillationCapability, +) +from lightx2v_train.model_zoo.wan.capability_adapters.wan_teacher_forcing_capability import ( + WanTeacherForcingCapability, +) from lightx2v_train.runtime.distributed import get_sequence_parallel_world_size from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype -from .base import BaseModel -from .native.wan.modules.causal_model import CausalWanModel -from .native.wan.modules.model import WanModel -from .native.wan.modules.t5 import T5EncoderModel -from .native.wan.modules.vae import WanVAE +from ..base import BaseModel +from ..native.wan.modules.causal_model import CausalWanModel +from ..native.wan.modules.model import WanModel +from ..native.wan.modules.t5 import T5EncoderModel +from ..native.wan.modules.vae import WanVAE @dataclass @@ -34,6 +50,26 @@ class WanT2VDenoiserInput: class WanT2VModel(BaseModel): pipeline_cls = None + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + FlowMatchingSFTCapability, + GenericFlowMatchingCapability(self), + ) + self.capabilities.register( + DistillationCapability, + WanDistillationCapability(self), + ) + if self.use_causal_transformer: + self.capabilities.register( + AutoregressiveDistillationCapability, + WanAutoregressiveDistillationCapability(self), + ) + self.capabilities.register( + TeacherForcingCapability, + WanTeacherForcingCapability(self), + ) + def load_components(self, transformer_only=False, reference_model=None): model_config = self.config["model"] model_path = model_config["pretrained_model_name_or_path"] @@ -249,6 +285,8 @@ def encode_to_latent(self, sample): latent = latent.to(device=self.device, dtype=self.running_dtype) if latent.ndim == 4: latent = latent.unsqueeze(0) + if latent.shape[0] != 1: + raise ValueError("Wan training only supports physical batch size 1.") return latent if self.vae is None: @@ -258,6 +296,8 @@ def encode_to_latent(self, sample): if video is None: raise KeyError("Wan encode_to_latent expects inputs.video or inputs.latents.") video = video.to(device=self.device, dtype=self.vae_dtype) + if video.shape[0] != 1: + raise ValueError("Wan training only supports physical batch size 1.") latent = torch.stack(self.vae.encode(self._batch_to_list(video)), dim=0) return latent.to(dtype=self.running_dtype) @@ -274,6 +314,8 @@ def encode_condition(self, sample): prompt_embed = prompt_embed.to(device=self.device, dtype=self.running_dtype) if prompt_embed.ndim == 2: prompt_embed = prompt_embed.unsqueeze(0) + if prompt_embed.shape[0] != 1: + raise ValueError("Wan cached conditions must have leading dimension 1.") return {"prompt_embed": prompt_embed} if self.text_encoder is None: @@ -281,6 +323,8 @@ def encode_condition(self, sample): prompt = conditioning.get("prompt", "") prompts = [prompt] if isinstance(prompt, str) else list(prompt) + if len(prompts) != 1: + raise ValueError(f"Wan training requires exactly one prompt, got {len(prompts)}.") text_device = torch.device("cpu") if self.t5_cpu else self.device contexts = self.text_encoder(prompts, text_device) prompt_embed = self._pad_contexts(contexts, device=self.device, dtype=self.running_dtype) @@ -301,6 +345,8 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): hidden_states = denoiser_input.hidden_states.to(device=self.device, dtype=self.running_dtype) if hidden_states.ndim == 4: hidden_states = hidden_states.unsqueeze(0) + if hidden_states.shape[0] != 1: + raise ValueError("Wan denoising only supports physical batch size 1 during training.") seq_len = self._sequence_length(hidden_states) if isinstance(self.transformer, CausalWanModel): @@ -352,6 +398,8 @@ def denoise_teacher_forcing( noisy_latent = noisy_latent.unsqueeze(0) if clean_latent.ndim == 4: clean_latent = clean_latent.unsqueeze(0) + if noisy_latent.shape[0] != 1 or clean_latent.shape[0] != 1: + raise ValueError("Wan teacher forcing only supports physical batch size 1.") timestep = self._causal_timestep(timestep, noisy_latent) if aug_t is not None: aug_t = self._causal_timestep(aug_t, noisy_latent) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_ti2v_5b.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py similarity index 95% rename from lightx2v_train/lightx2v_train/model_zoo/wan_ti2v_5b.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py index 2506567fc..cd90e97fb 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan_ti2v_5b.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py @@ -6,10 +6,10 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype -from .native.wan.modules.causal_model import CausalWanModel -from .native.wan.modules.model2_2 import WanModel as Wan2_2Model -from .native.wan.modules.t5 import T5EncoderModel -from .native.wan.modules.vae2_2 import Wan2_2_VAE +from ..native.wan.modules.causal_model import CausalWanModel +from ..native.wan.modules.model2_2 import WanModel as Wan2_2Model +from ..native.wan.modules.t5 import T5EncoderModel +from ..native.wan.modules.vae2_2 import Wan2_2_VAE from .wan_t2v import WanT2VModel @@ -130,8 +130,8 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): hidden_states = denoiser_input.hidden_states.to(device=self.device, dtype=self.running_dtype) if hidden_states.ndim == 4: hidden_states = hidden_states.unsqueeze(0) - if timestep.ndim == 1 and timestep.shape[0] == 1 and hidden_states.shape[0] > 1: - timestep = timestep.expand(hidden_states.shape[0]) + if hidden_states.shape[0] != 1: + raise ValueError("Wan TI2V training only supports physical batch size 1.") if isinstance(self.transformer, CausalWanModel): timestep = self._causal_timestep(timestep, hidden_states) diff --git a/lightx2v_train/lightx2v_train/schedulers/dmd_scheduler.py b/lightx2v_train/lightx2v_train/schedulers/dmd_scheduler.py index fa5a05e8e..772f4940b 100644 --- a/lightx2v_train/lightx2v_train/schedulers/dmd_scheduler.py +++ b/lightx2v_train/lightx2v_train/schedulers/dmd_scheduler.py @@ -71,15 +71,15 @@ def set_random_timesteps( self.timesteps = self.infer_timesteps self.num_inference_steps = int(sigmas.numel() - 1) - def sigma_at(self, step_idx, batch_size, device=None, dtype=None): - sigma = self.sigmas[int(step_idx)].expand(int(batch_size)) + def sigma_at(self, step_idx, device=None, dtype=None): + sigma = self.sigmas[int(step_idx)].reshape(1) if device is not None or dtype is not None: sigma = sigma.to(device=device, dtype=dtype) return sigma - def sample_renoise_sigma(self, batch_size, device=None, dtype=None): + def sample_renoise_sigma(self, device=None, dtype=None): device = device or self.device - raw = torch.rand((int(batch_size),), device=device, dtype=torch.float32) + raw = torch.rand((1,), device=device, dtype=torch.float32) if self.renoise_discrete_samples > 0: raw = torch.ceil(raw * self.renoise_discrete_samples) / self.renoise_discrete_samples raw = torch.clamp(raw, 1e-7, 1 - 1e-7) @@ -89,13 +89,17 @@ def sample_renoise_sigma(self, batch_size, device=None, dtype=None): return sigma def add_noise(self, latent, noise, sigmas): + if latent.shape[0] != 1 or noise.shape[0] != 1: + raise ValueError("DMD scheduler only supports physical batch size 1.") sigmas = sigmas.to(device=latent.device) sigmas = self._expand_to_ndim(sigmas, latent.ndim) return ((1.0 - sigmas) * latent + sigmas * noise).to(dtype=latent.dtype) def step_by_index(self, velocity, step_idx, sample): - sigma = self.sigma_at(step_idx, sample.shape[0], device=sample.device) - sigma_next = self.sigma_at(int(step_idx) + 1, sample.shape[0], device=sample.device) + if sample.shape[0] != 1: + raise ValueError("DMD scheduler only supports physical batch size 1.") + sigma = self.sigma_at(step_idx, device=sample.device) + sigma_next = self.sigma_at(int(step_idx) + 1, device=sample.device) sigma = self._expand_to_ndim(sigma, sample.ndim) sigma_next = self._expand_to_ndim(sigma_next, sample.ndim) next_sample = sample + (sigma_next - sigma) * velocity diff --git a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py index 7028ccbbe..958227e2c 100644 --- a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py +++ b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py @@ -43,27 +43,27 @@ def __init__(self, config): self.infer_timesteps = None self.num_inference_steps = None - def sample_timestep_or_sigma(self, num_samples, latent_hw=None, seq_len=None): + def sample_timestep_or_sigma(self, latent_hw=None, seq_len=None): if self.timestep_distribution == "logitnormal": - timestep_or_sigma = torch.randn((num_samples,), device=self.device, dtype=torch.float32) * self.logitnormal_std + self.logitnormal_mean + timestep_or_sigma = torch.randn((1,), device=self.device, dtype=torch.float32) * self.logitnormal_std + self.logitnormal_mean timestep_or_sigma = torch.sigmoid(timestep_or_sigma) timestep_or_sigma = timestep_or_sigma * (self.max_t - self.min_t) + self.min_t # [0, 1] -> [min_t, max_t] elif self.timestep_distribution == "uniform": - timestep_or_sigma = torch.rand((num_samples,), device=self.device) + timestep_or_sigma = torch.rand((1,), device=self.device) timestep_or_sigma = timestep_or_sigma * (self.max_t - self.min_t) + self.min_t # [0, 1] -> [min_t, max_t] elif self.timestep_distribution in {"shifted_logit_normal", "shifted_logitnormal"}: if seq_len is None: raise ValueError("scheduler.timestep_distribution='shifted_logit_normal' requires seq_len.") - timestep_or_sigma = self._sample_shifted_logit_normal(num_samples, seq_len) + timestep_or_sigma = self._sample_shifted_logit_normal(seq_len) else: raise ValueError(f"Unsupported timestep distribution: {self.timestep_distribution}") if self.do_time_shift: timestep_or_sigma = self.time_shift(timestep_or_sigma, latent_hw=latent_hw) return timestep_or_sigma.to(self.running_dtype) - def _sample_shifted_logit_normal(self, num_samples, seq_len): + def _sample_shifted_logit_normal(self, seq_len): mu = self._get_shift_for_sequence_length(seq_len) - normal = torch.randn((num_samples,), device=self.device, dtype=torch.float32) * self.logitnormal_std + mu + normal = torch.randn((1,), device=self.device, dtype=torch.float32) * self.logitnormal_std + mu samples = torch.sigmoid(normal) upper = torch.sigmoid(torch.tensor(mu + 3.0902 * self.logitnormal_std, device=self.device, dtype=torch.float32)) @@ -72,8 +72,8 @@ def _sample_shifted_logit_normal(self, num_samples, seq_len): stretched = torch.where(stretched >= self.logitnormal_eps, stretched, 2 * self.logitnormal_eps - stretched) stretched = torch.clamp(stretched, 0, 1) - uniform = (1 - self.logitnormal_eps) * torch.rand((num_samples,), device=self.device, dtype=torch.float32) + self.logitnormal_eps - choose_shifted = torch.rand((num_samples,), device=self.device) > self.logitnormal_uniform_prob + uniform = (1 - self.logitnormal_eps) * torch.rand((1,), device=self.device, dtype=torch.float32) + self.logitnormal_eps + choose_shifted = torch.rand((1,), device=self.device) > self.logitnormal_uniform_prob return torch.where(choose_shifted, stretched, uniform) @staticmethod @@ -175,33 +175,33 @@ def set_timesteps(self, num_inference_steps=1000, denoising_strength=1.0, traini y_shifted = y - y.min() self.linear_timesteps_weights = y_shifted * (num_inference_steps / y_shifted.sum()) - def sample_chunkwise(self, batch_size, num_frames, num_frame_per_chunk, device, dtype): + def sample_chunkwise(self, num_frames, num_frame_per_chunk, device, dtype): index = torch.randint( 0, self.num_train_timesteps, - (batch_size, num_frames), + (1, num_frames), device=device, dtype=torch.long, ) - index = index.reshape(batch_size, -1, num_frame_per_chunk) + index = index.reshape(1, -1, num_frame_per_chunk) index[:, :, 1:] = index[:, :, 0:1] - index = index.reshape(batch_size, num_frames) + index = index.reshape(1, num_frames) sigmas = self.sigmas.to(device=device, dtype=dtype)[index] weights = self.linear_timesteps_weights.to(device=device, dtype=torch.float32)[index] return sigmas, weights - def sample_clean_augmentation(self, batch_size, num_frames, num_frame_per_chunk, max_timestep, device, dtype): + def sample_clean_augmentation(self, num_frames, num_frame_per_chunk, max_timestep, device, dtype): index = torch.randint( int(max_timestep), self.num_train_timesteps, - (batch_size, num_frames), + (1, num_frames), device=device, dtype=torch.long, ) - index = index.reshape(batch_size, -1, num_frame_per_chunk) + index = index.reshape(1, -1, num_frame_per_chunk) index[:, :, 1:] = index[:, :, 0:1] - index = index.reshape(batch_size, num_frames) + index = index.reshape(1, num_frames) return self.sigmas.to(device=device, dtype=dtype)[index] def add_noise(self, latent, noise, sigmas): @@ -234,8 +234,8 @@ def _precompute_training_weight_stats(self): minimum = float(weights.min().item()) return minimum, float((weights - minimum).mean().item()) - def sample_training_t(self, batch_size, device, dtype): - value = torch.rand((batch_size,), device=device, dtype=torch.float32) + def sample_training_t(self, device, dtype): + value = torch.rand((1,), device=device, dtype=torch.float32) return (self._phi(value, self.shift) * self.num_train_timesteps).to(dtype=dtype) def training_weight(self, timestep): diff --git a/lightx2v_train/lightx2v_train/trainers/__init__.py b/lightx2v_train/lightx2v_train/trainers/__init__.py index cdae33df5..225a8cab4 100644 --- a/lightx2v_train/lightx2v_train/trainers/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/__init__.py @@ -3,43 +3,30 @@ from lightx2v_train.utils.registry import build_trainer _LAZY_EXPORTS = { - "ARDmdTrainer": (".dmd.video_ar_trainer", "VideoArDmdTrainer"), - "ConsistencyTrainer": (".consistency.trainer", "ConsistencyTrainer"), - "DmdTrainer": (".dmd.trainer", "DmdTrainer"), - "DopsdTrainer": (".dopsd", "DopsdTrainer"), - "FastWAMTrainer": (".fastwam", "FastWAMTrainer"), - "FlowMatchingTrainer": (".flow", "FlowMatchingTrainer"), - "LTX2T2AVArDmdTrainer": ( - ".dmd.ltx_trainer", - "LTX2T2AVArDmdTrainer", - ), - "LTX2T2AVDmdTrainer": ( - ".dmd.ltx_trainer", - "LTX2T2AVDmdTrainer", - ), - "LTX2T2AVFlowTrainer": (".flow", "LTX2T2AVFlowTrainer"), - "LTX2T2AVTeacherForcingTrainer": ( - ".tf", - "LTX2T2AVTeacherForcingTrainer", + "AutoregressiveDmdTrainer": ( + ".dmd.autoregressive_dmd", + "AutoregressiveDmdTrainer", ), - "LingBotVideoDmdTrainer": ( - ".dmd.video_trainer", - "LingBotVideoDmdTrainer", + "ConsistencyTrainer": ( + ".consistency.trainer", + "ConsistencyTrainer", ), - "TFTrainer": (".tf", "TFTrainer"), - "VideoArDmdTrainer": ( - ".dmd.video_ar_trainer", - "VideoArDmdTrainer", - ), - "VideoDmdTrainer": ( - ".dmd.video_trainer", - "VideoDmdTrainer", + "DmdTrainer": (".dmd.trainer", "DmdTrainer"), + "DopsdTrainer": (".dopsd", "DopsdTrainer"), + "FlowMatchingTrainer": ( + ".flow_matching", + "FlowMatchingTrainer", ), - "VideoPhasedDmdTrainer": ( + "PhasedDmdTrainer": ( ".phased_dmd.trainer", - "VideoPhasedDmdTrainer", + "PhasedDmdTrainer", + ), + "SgmdTrainer": (".sgmd", "SgmdTrainer"), + "TeacherForcingTrainer": ( + ".teacher_forcing", + "TeacherForcingTrainer", ), - "VideoSgmdTrainer": (".sgmd", "VideoSgmdTrainer"), + "WorldActionTrainer": (".world_action", "WorldActionTrainer"), } @@ -48,27 +35,12 @@ def __getattr__(name): if target is None: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") module_name, attribute_name = target - value = getattr(importlib.import_module(module_name, __name__), attribute_name) + value = getattr( + importlib.import_module(module_name, __name__), + attribute_name, + ) globals()[name] = value return value -__all__ = [ - "build_trainer", - "ARDmdTrainer", - "ConsistencyTrainer", - "DmdTrainer", - "FlowMatchingTrainer", - "LTX2T2AVArDmdTrainer", - "LTX2T2AVDmdTrainer", - "LTX2T2AVFlowTrainer", - "LTX2T2AVTeacherForcingTrainer", - "LingBotVideoDmdTrainer", - "TFTrainer", - "VideoArDmdTrainer", - "VideoDmdTrainer", - "VideoPhasedDmdTrainer", - "VideoSgmdTrainer", - "DopsdTrainer", - "FastWAMTrainer", -] +__all__ = ["build_trainer", *_LAZY_EXPORTS] diff --git a/lightx2v_train/lightx2v_train/trainers/base.py b/lightx2v_train/lightx2v_train/trainers/base.py index 4234ec9ef..62f639b7a 100644 --- a/lightx2v_train/lightx2v_train/trainers/base.py +++ b/lightx2v_train/lightx2v_train/trainers/base.py @@ -7,16 +7,28 @@ from loguru import logger from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict +from lightx2v_train.data.utils import require_singleton_dataloader from lightx2v_train.infer import build_inferencer +from lightx2v_train.model_capabilities import ( + CapabilityNotSupportedError, + CheckpointCapability, + ParallelCapability, + TrainableModelCapability, +) from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints from lightx2v_train.runtime.distributed import barrier, get_world_size, is_main_process from lightx2v_train.runtime.monitor import build_monitor -from lightx2v_train.runtime.parallel import apply_parallel, set_parallel_gradient_sync from lightx2v_train.schedulers.flow_matching import RectifiedFlowMatchingScheduler from lightx2v_train.utils.utils import get_running_dtype class BaseTrainer: + required_capabilities = ( + TrainableModelCapability, + ParallelCapability, + CheckpointCapability, + ) + def __init__(self, config): self.config = config self.model_config = self.config["model"] @@ -81,8 +93,23 @@ def _get_optimizer_config(self): def set_model(self, model): self.model = model + registry = model.ensure_capabilities() + missing = registry.missing(self.capability_requirements()) + if missing: + model_name = self.model_config.get("name", type(model).__name__) + names = ", ".join(capability.__name__ for capability in missing) + raise CapabilityNotSupportedError(f"trainer={type(self).__name__} model={model_name!r} is missing capabilities: {names}.") + self.trainable_model = registry.require(TrainableModelCapability) + self.parallel = registry.require(ParallelCapability) + self.checkpoint = registry.require(CheckpointCapability) + + def capability_requirements(self): + return self.required_capabilities def set_data(self, dataloader_train, dataloader_eval=None): + require_singleton_dataloader(dataloader_train, "Training dataloader") + if dataloader_eval is not None: + require_singleton_dataloader(dataloader_eval, "Evaluation dataloader") self.dataloader_train = dataloader_train self.dataloader_eval = dataloader_eval @@ -93,17 +120,18 @@ def finish_monitor(self): self.monitor.finish() def _setup_trainable_model(self, model): - if self.train_type == "lora": - model.add_lora(self.lora_rank, self.lora_alpha, self.lora_target_modules) - model.set_lora_trainable() - return - model.set_full_trainable() + capability = model.ensure_capabilities().require(TrainableModelCapability) + capability.configure( + self.train_type, + { + "rank": self.lora_rank, + "alpha": self.lora_alpha, + "target_modules": self.lora_target_modules, + }, + ) def _restore_trainable_model(self, model): - if self.train_type == "lora": - model.set_lora_trainable() - return - model.set_full_trainable() + model.ensure_capabilities().require(TrainableModelCapability).restore(self.train_type) def _build_optimizer(self, params, optimizer_config=None): if optimizer_config is None: @@ -133,18 +161,18 @@ def _build_lr_scheduler(self, optimizer, num_training_steps=None, num_warmup_ste def setup(self, resume_ckpt_path=None): self._setup_trainable_model(self.model) - apply_parallel(self.model, self.config) + self.parallel.apply(self.config) if self.gradient_checkpointing: - self.model.enable_gradient_checkpointing() + self.trainable_model.enable_gradient_checkpointing() if self.infer_every_iters: self.inferencer = build_inferencer(self.config) self.inferencer.set_model(self.model) - self.model.log_model_structure() + self.trainable_model.log_structure() - self.trainable_params = list(self.model.trainable_parameters()) + self.trainable_params = list(self.trainable_model.parameters()) self.optimizer = self._build_optimizer(self.trainable_params) self.lr_scheduler = self._build_lr_scheduler(self.optimizer) @@ -152,25 +180,13 @@ def setup(self, resume_ckpt_path=None): self._load_resume_state(resume_ckpt_path) def _save_model_weights(self, model, save_dir): - if self.train_type == "lora": - model.save_lora_weights(save_dir) - return - if is_main_process(): - torch.save(model.denoiser_module().state_dict(), os.path.join(save_dir, "model_state.pt")) + model.ensure_capabilities().require(CheckpointCapability).save_weights(save_dir, self.train_type) def _load_model_weights(self, model, save_dir): - if self.train_type == "lora": - model.load_lora_weights_for_resume(save_dir) - model.load_consistency_auxiliary_weights(save_dir) - return - model_state_path = os.path.join(save_dir, "model_state.pt") - if not os.path.exists(model_state_path): - raise RuntimeError(f"model_state.pt not found in {save_dir}") - state_dict = torch.load(model_state_path, map_location="cpu", weights_only=False) - model.denoiser_module().load_state_dict(state_dict) + model.ensure_capabilities().require(CheckpointCapability).load_weights(save_dir, self.train_type) def _load_resume_state(self, resume_ckpt_path): - if self.model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp(): self._load_distributed_state(resume_ckpt_path) return @@ -200,11 +216,16 @@ def _load_distributed_state(self, resume_ckpt_path): self._validate_checkpoint_metadata(trainer_state, trainer_state_path, resume_ckpt_path) options = StateDictOptions(ignore_frozen_params=True, strict=False) - model_state, optim_state = get_state_dict(self.model.fsdp2_state_module(), self.optimizer, options=options) + state_module = self.parallel.state_module() + model_state, optim_state = get_state_dict( + state_module, + self.optimizer, + options=options, + ) state = {"model": model_state, "optimizer": optim_state} dcp.load(state, checkpoint_id=dist_state_path) set_state_dict( - self.model.fsdp2_state_module(), + state_module, self.optimizer, model_state_dict=state["model"], optim_state_dict=state["optimizer"], @@ -236,7 +257,7 @@ def _resolve_resume(self): return ckpt_path, current_iter def _set_gradient_sync(self, enabled): - set_parallel_gradient_sync(self.model, enabled) + self.parallel.set_gradient_sync(enabled) def _after_backward(self): pass @@ -267,7 +288,7 @@ def save_checkpoint(self, iteration, save_total_limit): os.makedirs(save_dir, exist_ok=True) barrier() - save_standalone_weights = self.train_type == "lora" or not self.model.is_fsdp2_wrapped() + save_standalone_weights = self.train_type == "lora" or not self.parallel.is_fsdp() if save_standalone_weights: self._save_model_weights(self.model, save_dir) barrier() @@ -276,7 +297,7 @@ def save_checkpoint(self, iteration, save_total_limit): if is_main_process() and config_path is not None: shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) - if self.model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp(): self._save_distributed_state(save_dir, iteration) self._save_consolidated_weights(save_dir) barrier() @@ -302,7 +323,7 @@ def _save_consolidated_weights(self, save_dir): logger.warning("training.save_consolidated_weights=true is ignored for train_type='{}'.", self.train_type) return output_path = os.path.join(save_dir, self.consolidated_weights_name) - self.model.save_consolidated_weights(output_path) + self.checkpoint.save_consolidated(output_path) barrier() def _save_distributed_state(self, save_dir, iteration): @@ -320,7 +341,11 @@ def _save_distributed_state(self, save_dir, iteration): barrier() options = StateDictOptions(ignore_frozen_params=True, strict=False) - model_state, optim_state = get_state_dict(self.model.fsdp2_state_module(), self.optimizer, options=options) + model_state, optim_state = get_state_dict( + self.parallel.state_module(), + self.optimizer, + options=options, + ) dcp.save( {"model": model_state, "optimizer": optim_state}, checkpoint_id=dist_state_path, diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py b/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py index ad763f133..d99bf5773 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/__init__.py @@ -1,9 +1,9 @@ from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyObjective, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, ReferenceModelSpec, @@ -16,7 +16,7 @@ "ConsistencyObjective", "ConsistencyStepContext", "DenoiserRequest", - "ModelDenoiser", + "CapabilityDenoiser", "ObjectiveOutput", "RectifiedFlowPath", "ReferenceModelSpec", diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/base.py b/lightx2v_train/lightx2v_train/trainers/consistency/base.py index 413688b04..3b4d49cd4 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/base.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/base.py @@ -32,6 +32,11 @@ def expand_time(time: Tensor, ndim: int) -> Tensor: return time.reshape(time.shape[0], *([1] * (ndim - 1))) +def require_singleton_clean(clean: Tensor) -> None: + if clean.ndim == 0 or clean.shape[0] != 1: + raise ValueError("Consistency training only supports physical batch size 1.") + + class RectifiedFlowPath: """Conversions for the straight path x_t=(1-t)x_0+t*noise.""" @@ -87,34 +92,19 @@ class DenoiserRequest: model_kwargs: Mapping[str, Any] = field(default_factory=dict) -class ModelDenoiser: - """Adapt a LightX2V model to the consistency objective interface.""" +class CapabilityDenoiser: + """Adapt a consistency capability to the objective interface.""" - def __init__(self, model, path: RectifiedFlowPath): - self.model = model + def __init__(self, capability, path: RectifiedFlowPath): + self.capability = capability self.path = path def predict(self, request: DenoiserRequest) -> Tensor: - raw_prediction = self.model.predict_denoiser_output( - request.sample, - request.time, - request.condition, - **request.model_kwargs, - ) - return self.path.convert_prediction( - request.sample, - raw_prediction, - request.time, - source_type=self.model.denoiser_prediction_type(), - target_type=request.prediction_type, - ) + return self.capability.predict(request, self.path) def predict_log_variance(self, time: Tensor) -> Tensor: """Return the model-owned scalar log-variance head used by sCM.""" - predictor = getattr(self.model, "predict_consistency_log_variance", None) - if predictor is None: - raise NotImplementedError(f"{self.model.__class__.__name__} does not provide a consistency log-variance head.") - return predictor(time) + return self.capability.predict_log_variance(time) @dataclass(frozen=True) @@ -181,8 +171,8 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: """Compute a differentiable scalar loss and detached metrics.""" diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/cm.py b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py index fd57e8a37..33cdf4edc 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/cm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py @@ -7,14 +7,15 @@ from torch import Tensor from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyObjective, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, expand_time, + require_singleton_clean, ) from .config import CMConfig, CMLossConfig, CMTimePairConfig from .objective_factory import CONSISTENCY_OBJECTIVE_REGISTER @@ -160,9 +161,9 @@ def sample_training_state( scheduler, context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: - batch_size = clean.shape[0] + require_singleton_clean(clean) latent_hw = (clean.shape[-2], clean.shape[-1]) - sampled_t = scheduler.sample_timestep_or_sigma(batch_size, latent_hw=latent_hw).to(clean.device) + sampled_t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device) pair = self.time_pair_sampler.sample(sampled_t, context) return { "noise": torch.randn_like(clean), @@ -176,9 +177,9 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: del references terms = self.compute_loss_terms(batch, training_state, student, teacher) @@ -201,8 +202,8 @@ def compute_loss_terms( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, ) -> CMLossTerms: clean = batch.clean noise = training_state["noise"] @@ -269,7 +270,7 @@ def _distill_to_r( r: Tensor, condition, negative_condition, - teacher: ModelDenoiser, + teacher: CapabilityDenoiser, ) -> Tensor: conditional_velocity = teacher.predict( DenoiserRequest( diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py index ad5c36418..dea5a345e 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py @@ -6,14 +6,15 @@ from torch import Tensor from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyObjective, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, expand_time, + require_singleton_clean, ) from .config import MeanFlowConfig from .jvp import math_attention_for_forward_ad @@ -57,12 +58,13 @@ def sample_training_state( scheduler, context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: + require_singleton_clean(clean) latent_hw = (clean.shape[-2], clean.shape[-1]) - first = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() - second = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() + first = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() + second = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() t = torch.maximum(first, second) random_r = torch.minimum(first, second) - use_random_r = torch.rand(clean.shape[0], device=clean.device) < self.config.random_endpoint_probability + use_random_r = torch.rand(1, device=clean.device) < self.config.random_endpoint_probability r = torch.where(use_random_r, random_r, t) return { "noise": torch.randn_like(clean), @@ -76,9 +78,9 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: del references clean = batch.clean @@ -129,8 +131,8 @@ def _target_velocity( x_t: Tensor, noise: Tensor, t: Tensor, - student: ModelDenoiser, - teacher: Optional[ModelDenoiser], + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser], ): if self.config.mode == "cd": if teacher is None: @@ -196,7 +198,7 @@ def _jvp( r: Tensor, dxt_dt: Tensor, condition, - student: ModelDenoiser, + student: CapabilityDenoiser, ) -> Tensor: def model_fn(sample, time, endpoint): return student.predict( diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py index ea21cb6ed..3bcbebc41 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -7,13 +7,14 @@ from torch import Tensor from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyObjective, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, + require_singleton_clean, ) from .cm import classifier_free_guidance from .config import PCMConfig, PCMLossConfig, PCMSolverConfig @@ -41,9 +42,9 @@ class PCMTimeGrid: def __init__(self, config: PCMSolverConfig): self.config = config - def sample(self, batch_size: int, scheduler, *, latent_hw, device) -> PCMTimeState: + def sample(self, scheduler, *, latent_hw, device) -> PCMTimeState: current, previous = self.build_grid(scheduler, latent_hw=latent_hw, device=device) - solver_index = torch.randint(0, self.config.num_solver_steps, (batch_size,), device=device) + solver_index = torch.randint(0, self.config.num_solver_steps, (1,), device=device) phase_starts = self.phase_start_indices(device) phase_index = torch.bucketize(solver_index, phase_starts, right=True) - 1 phase_start_index = phase_starts[phase_index] @@ -133,9 +134,9 @@ def sample_training_state( context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: del context + require_singleton_clean(clean) latent_hw = (clean.shape[-2], clean.shape[-1]) time = self.time_grid.sample( - clean.shape[0], scheduler, latent_hw=latent_hw, device=clean.device, @@ -154,9 +155,9 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: del references if teacher is None: @@ -209,7 +210,7 @@ def _teacher_velocity( time: Tensor, condition, negative_condition, - teacher: ModelDenoiser, + teacher: CapabilityDenoiser, ) -> Tensor: conditional = teacher.predict(DenoiserRequest(sample, time, condition, prediction_type="velocity")) scale = self.config.teacher.guidance_scale diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/scm.py b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py index ade52e30f..cd9e1e69b 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/scm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py @@ -6,14 +6,15 @@ from torch import Tensor from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyObjective, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, expand_time, + require_singleton_clean, ) from .cm import classifier_free_guidance from .config import SCMConfig @@ -25,7 +26,7 @@ class SCMObjective(ConsistencyObjective): """Continuous-time simplified Consistency Model (sCM) objective. - The public model still predicts Qwen's rectified-flow velocity. This + The public model still predicts rectified-flow velocity. This objective performs FastGen's SNR-matched TrigFlow preconditioning at the model boundary, keeping the backbone and its checkpoints interoperable. """ @@ -47,8 +48,9 @@ def sample_training_state( scheduler, context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: + require_singleton_clean(clean) latent_hw = (clean.shape[-2], clean.shape[-1]) - t = scheduler.sample_timestep_or_sigma(clean.shape[0], latent_hw=latent_hw).to(clean.device).float() + t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() return { "noise": torch.randn_like(clean) * self.config.sigma_data, "t": t, @@ -59,9 +61,9 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: del references clean = batch.clean @@ -130,7 +132,7 @@ def compute( }, ) - def _predict_trig_flow(self, denoiser: ModelDenoiser, x_hat: Tensor, t_hat: Tensor, condition): + def _predict_trig_flow(self, denoiser: CapabilityDenoiser, x_hat: Tensor, t_hat: Tensor, condition): sigma_data = self.config.sigma_data tangent = sigma_data * torch.tan(t_hat.double()) original_t = (tangent / (1.0 + tangent)).to(t_hat.dtype) @@ -157,7 +159,7 @@ def _jvp( t_hat: Tensor, dxt_dt: Tensor, condition, - student: ModelDenoiser, + student: CapabilityDenoiser, ) -> Tensor: def model_fn(sample, time): return self._predict_trig_flow(student, sample, time, condition)[0] diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py index 21b984523..66a4b8fcd 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py @@ -7,14 +7,15 @@ from torch import Tensor from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyStepContext, DenoiserRequest, - ModelDenoiser, ObjectiveOutput, RectifiedFlowPath, ReferenceModelSpec, expand_time, + require_singleton_clean, ) from .cm import CMObjective from .config import TCMConfig @@ -34,7 +35,7 @@ def _distributed_mask_summary(mask: Tensor) -> tuple[bool, bool]: class TwoStageDenoiser: """Route low times through frozen stage 1 and high times through stage 2.""" - def __init__(self, stage1: ModelDenoiser, stage2: ModelDenoiser, transition_time: float): + def __init__(self, stage1: CapabilityDenoiser, stage2: CapabilityDenoiser, transition_time: float): self.stage1 = stage1 self.stage2 = stage2 self.transition_time = transition_time @@ -88,18 +89,12 @@ def sample_training_state( scheduler, context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: - batch_size = clean.shape[0] + require_singleton_clean(clean) latent_hw = (clean.shape[-2], clean.shape[-1]) - t = scheduler.sample_timestep_or_sigma(batch_size, latent_hw=latent_hw).to(clean.device).float() + t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() t = t.clamp_min(self.tcm_config.transition_time + self.config.time_pair.safety_epsilon) - boundary_count = int(batch_size * self.tcm_config.boundary_probability) - if boundary_count: - permutation = torch.randperm(batch_size, device=clean.device) - boundary_mask = torch.zeros(batch_size, device=clean.device, dtype=torch.bool) - boundary_mask[permutation[:boundary_count]] = True - else: - boundary_mask = torch.rand(batch_size, device=clean.device) < self.tcm_config.boundary_probability + boundary_mask = torch.rand(1, device=clean.device) < self.tcm_config.boundary_probability t = torch.where( boundary_mask, torch.full_like(t, self.tcm_config.transition_time + self.config.time_pair.safety_epsilon), @@ -120,9 +115,9 @@ def compute( self, batch: ConsistencyBatch, training_state: Mapping[str, Tensor], - student: ModelDenoiser, - teacher: Optional[ModelDenoiser] = None, - references: Optional[Mapping[str, ModelDenoiser]] = None, + student: CapabilityDenoiser, + teacher: Optional[CapabilityDenoiser] = None, + references: Optional[Mapping[str, CapabilityDenoiser]] = None, ) -> ObjectiveOutput: if references is None or "stage1" not in references: raise RuntimeError("TCM requires the frozen 'stage1' consistency model.") diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py index d3c1bd4f9..8151ba3af 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py @@ -6,23 +6,28 @@ import torch from loguru import logger -from lightx2v_train.model_zoo import build_model +from lightx2v_train.model_capabilities import ( + ConsistencyCapability, + LossResult, + ParallelCapability, + TrainableModelCapability, +) +from lightx2v_train.model_zoo import build_loaded_model from lightx2v_train.runtime.distributed import get_data_parallel_world_size -from lightx2v_train.runtime.parallel import apply_parallel from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value from lightx2v_train.utils.registry import TRAINER_REGISTER -from ..flow import FlowMatchingTrainer +from ..base import BaseTrainer +from ..flow_matching import FlowMatchingTrainer from .base import ( + CapabilityDenoiser, ConsistencyBatch, ConsistencyStepContext, - ModelDenoiser, RectifiedFlowPath, ) from .objective_factory import build_consistency_objective -@TRAINER_REGISTER("cm") @TRAINER_REGISTER("consistency") class ConsistencyTrainer(FlowMatchingTrainer): """Shared trainer shell for consistency-model objectives. @@ -33,6 +38,10 @@ class ConsistencyTrainer(FlowMatchingTrainer): """ trainer_name = "consistency" + required_capabilities = ( + *BaseTrainer.required_capabilities, + ConsistencyCapability, + ) def __init__(self, config): super().__init__(config) @@ -44,19 +53,32 @@ def __init__(self, config): self.reference_denoisers = {} self._setup_resume_checkpoint = None + def set_model(self, model): + BaseTrainer.set_model(self, model) + self.consistency = model.capabilities.require(ConsistencyCapability) + def setup(self, resume_ckpt_path=None): self._setup_resume_checkpoint = resume_ckpt_path super().setup(resume_ckpt_path=resume_ckpt_path) - self.student_denoiser = ModelDenoiser(self.model, self.path) + self.student_denoiser = CapabilityDenoiser( + self.consistency, + self.path, + ) if self.objective.requires_teacher: self.teacher_model = self._build_frozen_teacher() - self.teacher_denoiser = ModelDenoiser(self.teacher_model, self.path) + self.teacher_denoiser = CapabilityDenoiser( + self.teacher_consistency, + self.path, + ) for spec in self.objective.reference_model_specs: reference_model = self._build_frozen_reference(spec) self.reference_models[spec.role] = reference_model - self.reference_denoisers[spec.role] = ModelDenoiser(reference_model, self.path) + self.reference_denoisers[spec.role] = CapabilityDenoiser( + self.reference_capabilities[spec.role], + self.path, + ) logger.info( "[train] consistency algorithm={} mode={} teacher={}", @@ -66,20 +88,21 @@ def setup(self, resume_ckpt_path=None): ) def _setup_trainable_model(self, model): - model.configure_consistency_model(self.objective.model_capabilities) + capability = model.ensure_capabilities().require(ConsistencyCapability) + capability.configure(self.objective.model_capabilities) super()._setup_trainable_model(model) - model.set_consistency_modules_trainable() + capability.restore_trainable_auxiliary() initialization = self.objective.student_initialization_checkpoint if initialization is not None and self._setup_resume_checkpoint is None: self._load_initial_model_weights(model, initialization) # Loading weights does not change requires_grad, but keeping this # call here makes that lifecycle guarantee explicit for new models. - model.set_consistency_modules_trainable() + capability.restore_trainable_auxiliary() def _restore_trainable_model(self, model): super()._restore_trainable_model(model) - model.set_consistency_modules_trainable() + model.ensure_capabilities().require(ConsistencyCapability).restore_trainable_auxiliary() def _load_initial_model_weights(self, model, checkpoint, *, role="student"): if not os.path.isdir(checkpoint): @@ -109,12 +132,15 @@ def _build_frozen_teacher(self): teacher_config = copy.deepcopy(self.config) teacher_config["model"] = base_model_config teacher_config["model"].update(copy.deepcopy(teacher_override)) - teacher_model = build_model(teacher_config) - teacher_model.load_components(transformer_only=True, reference_model=self.model) - teacher_model.denoiser_module().requires_grad_(False) - teacher_model.set_denoiser_eval() - apply_parallel(teacher_model, self.config) - teacher_model.set_denoiser_eval() + teacher_model = build_loaded_model( + teacher_config, + transformer_only=True, + reference_model=self.model, + ) + self.teacher_consistency = teacher_model.capabilities.require(ConsistencyCapability) + self.teacher_consistency.set_frozen() + teacher_model.capabilities.require(ParallelCapability).apply(self.config) + self.teacher_consistency.set_frozen() logger.info( "[train] consistency teacher model={} path={}", teacher_config["model"]["name"], @@ -144,21 +170,28 @@ def _build_frozen_reference(self, spec): reference_config = copy.deepcopy(self.config) reference_config["model"] = model_config - reference_model = build_model(reference_config) - reference_model.load_components(transformer_only=True, reference_model=self.model) + reference_model = build_loaded_model( + reference_config, + transformer_only=True, + reference_model=self.model, + ) + capability = reference_model.capabilities.require(ConsistencyCapability) if self.train_type == "lora": - reference_model.add_lora(self.lora_rank, self.lora_alpha, self.lora_target_modules) + reference_model.capabilities.require(TrainableModelCapability).configure( + "lora", + { + "rank": self.lora_rank, + "alpha": self.lora_alpha, + "target_modules": self.lora_target_modules, + }, + ) self._load_initial_model_weights(reference_model, spec.checkpoint, role=spec.role) - reference_model.denoiser_module().requires_grad_(False) - if spec.training_mode: - reference_model.denoiser_module().train() - else: - reference_model.set_denoiser_eval() - apply_parallel(reference_model, self.config) - if spec.training_mode: - reference_model.denoiser_module().train() - else: - reference_model.set_denoiser_eval() + capability.set_frozen(training=spec.training_mode) + reference_model.capabilities.require(ParallelCapability).apply(self.config) + capability.set_frozen(training=spec.training_mode) + if not hasattr(self, "reference_capabilities"): + self.reference_capabilities = {} + self.reference_capabilities[spec.role] = capability logger.info( "[train] consistency reference role={} model={} checkpoint={}", spec.role, @@ -169,14 +202,14 @@ def _build_frozen_reference(self, spec): def compute_loss_on_sample(self, sample): with torch.no_grad(): - clean = self.model.encode_to_latent(sample) + clean = self.consistency.encode_latent(sample) clean = broadcast_sequence_parallel_value(clean) - condition = self.model.encode_condition(sample) + condition = self.consistency.encode_condition(sample) condition = broadcast_sequence_parallel_value(condition) negative_condition = None if self.objective.requires_negative_condition: - negative_condition = self._encode_negative_condition(sample, clean.shape[0]) + negative_condition = self._encode_negative_condition(sample) negative_condition = broadcast_sequence_parallel_value(negative_condition) context = ConsistencyStepContext( @@ -201,29 +234,27 @@ def compute_loss_on_sample(self, sample): self.teacher_denoiser, self.reference_denoisers, ) - return {"loss": output.loss, "metrics": output.metrics} + return LossResult(loss=output.loss, metrics=output.metrics) - def _encode_negative_condition(self, sample, batch_size): + def _encode_negative_condition(self, sample): conditioning = sample.get("conditioning", {}) prompt = conditioning.get("prompt", "") negative_prompt = conditioning.get("negative_prompt") if negative_prompt is None: - values = [self.objective.negative_prompt] * batch_size + values = [self.objective.negative_prompt] elif isinstance(negative_prompt, str): - values = [negative_prompt] * batch_size + values = [negative_prompt] else: values = list(negative_prompt) - if len(values) == 1 and batch_size > 1: - values *= batch_size - elif len(values) != batch_size: - raise ValueError(f"Expected {batch_size} negative prompts, got {len(values)}.") + if len(values) != 1: + raise ValueError(f"Expected exactly one negative prompt, got {len(values)}.") fallback = self.objective.negative_prompt or " " values = [value if isinstance(value, str) and value.strip() else fallback for value in values] - encoded_prompt = values[0] if isinstance(prompt, str) and batch_size == 1 else values + encoded_prompt = values[0] if isinstance(prompt, str) else values negative_sample = dict(sample) negative_sample["conditioning"] = dict(conditioning) negative_sample["conditioning"]["prompt"] = encoded_prompt - return self.model.encode_condition(negative_sample) + return self.consistency.encode_condition(negative_sample) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/__init__.py b/lightx2v_train/lightx2v_train/trainers/dmd/__init__.py index d957f5958..7947efc98 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/__init__.py @@ -1,25 +1,11 @@ import importlib _LAZY_EXPORTS = { - "DmdTrainer": (".trainer", "DmdTrainer"), - "VideoDmdTrainer": (".video_trainer", "VideoDmdTrainer"), - "LingBotVideoDmdTrainer": ( - ".video_trainer", - "LingBotVideoDmdTrainer", - ), - "VideoArDmdTrainer": ( - ".video_ar_trainer", - "VideoArDmdTrainer", - ), - "ARDmdTrainer": (".video_ar_trainer", "VideoArDmdTrainer"), - "LTX2T2AVDmdTrainer": ( - ".ltx_trainer", - "LTX2T2AVDmdTrainer", - ), - "LTX2T2AVArDmdTrainer": ( - ".ltx_trainer", - "LTX2T2AVArDmdTrainer", + "AutoregressiveDmdTrainer": ( + ".autoregressive_dmd", + "AutoregressiveDmdTrainer", ), + "DmdTrainer": (".trainer", "DmdTrainer"), } diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py new file mode 100644 index 000000000..582a841f5 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py @@ -0,0 +1,72 @@ +from lightx2v_train.model_capabilities import ( + AutoregressiveDistillationCapability, + AutoregressiveRolloutContext, +) +from lightx2v_train.utils.registry import TRAINER_REGISTER + +from .trainer import DmdTrainer + + +@TRAINER_REGISTER("autoregressive_dmd") +class AutoregressiveDmdTrainer(DmdTrainer): + trainer_name = "autoregressive_dmd" + required_capabilities = ( + *DmdTrainer.required_capabilities, + AutoregressiveDistillationCapability, + ) + supports_diversity_loss = False + supports_real_data_fake = False + + def __init__(self, config): + super().__init__(config) + self.num_frame_per_chunk = int( + self.dmd_config.get( + "num_frame_per_chunk", + self.model_config.get("num_frame_per_chunk", 3), + ) + ) + self.same_step_across_blocks = bool(self.dmd_config.get("same_step_across_blocks", True)) + self.context_noise = float(self.dmd_config.get("context_noise", 0.0)) + self.sequence_parallel_cache = bool(self.dmd_config.get("sp_cache", False)) + + def set_model(self, model): + super().set_model(model) + self.autoregressive = model.capabilities.require(AutoregressiveDistillationCapability) + + def run_back_simulation( + self, + condition, + latent_shape, + grad_enabled, + xt=None, + ): + if xt is None: + xt = self.sample_initial_latents(latent_shape) + self.scheduler.set_timesteps( + self.num_inference_steps, + sigmas=[float(sigma) for sigma in self.denoising_sigmas.detach().cpu()], + latent_hw=self.student.latent_hw(latent_shape), + device=self.student.device, + ) + generated, exit_indices = self.autoregressive.rollout( + condition, + latent_shape, + xt, + AutoregressiveRolloutContext( + denoising_steps=self.denoising_steps, + denoising_scheduler=self.denoising_scheduler, + trajectory_scheduler=self.scheduler, + running_dtype=self.running_dtype, + frames_per_chunk=self.num_frame_per_chunk, + same_step_across_blocks=self.same_step_across_blocks, + context_noise=self.context_noise, + sequence_parallel_cache=self.sequence_parallel_cache, + grad_enabled=grad_enabled, + ), + ) + return generated, *self._denoised_timestep_window(exit_indices) + + def _denoised_timestep_window(self, exit_indices): + if not self.same_step_across_blocks: + return None, None + return super()._denoised_timestep_window(int(exit_indices[0])) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/checkpoint.py b/lightx2v_train/lightx2v_train/trainers/dmd/checkpoint.py index 1b7ae15b3..fbc5bad36 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/checkpoint.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/checkpoint.py @@ -13,6 +13,10 @@ set_state_dict, ) +from lightx2v_train.model_capabilities import ( + CheckpointCapability, + ParallelCapability, +) from lightx2v_train.runtime.checkpoint import prune_checkpoints from lightx2v_train.runtime.distributed import ( barrier, @@ -37,12 +41,22 @@ def _fake_weights_dir(self, root_dir): directory_name = self.role_registry.weight_directory_name("fake") return os.path.join(root_dir, directory_name) + @staticmethod + def _parallel(model): + return model.ensure_capabilities().require(ParallelCapability) + + @staticmethod + def _checkpoint(model): + return model.ensure_capabilities().require(CheckpointCapability) + def _copy_role_model(self, source_role, target_role): source = self.role_registry.runtime(source_role).model target = self.role_registry.runtime(target_role).model - source_module = source.fsdp2_state_module() - target_module = target.fsdp2_state_module() - if source.is_fsdp2_wrapped(): + source_parallel = self._parallel(source) + target_parallel = self._parallel(target) + source_module = source_parallel.state_module() + target_module = target_parallel.state_module() + if source_parallel.is_fsdp(): options = StateDictOptions( ignore_frozen_params=False, strict=True, @@ -132,7 +146,7 @@ def _validate_optional_trick_metadata(self, state, state_path): raise RuntimeError(f"Checkpoint {key}={state[key]!r} does not match the current value {value!r}: {state_path}") def _load_resume_state(self, resume_ckpt_path): - if self.model.is_fsdp2_wrapped() or self.fake_model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp() or self._parallel(self.fake_model).is_fsdp(): self._load_distributed_state(resume_ckpt_path) return @@ -280,8 +294,10 @@ def _load_distributed_state(self, resume_ckpt_path): self._validate_dmd_checkpoint_metadata(trainer_state, trainer_state_path, resume_ckpt_path) options = StateDictOptions(ignore_frozen_params=True, strict=False) - student_model_state, student_optim_state = get_state_dict(self.model.fsdp2_state_module(), self.optimizer, options=options) - fake_model_state, fake_optim_state = get_state_dict(self.fake_model.fsdp2_state_module(), self.fake_optimizer, options=options) + student_module = self.parallel.state_module() + fake_module = self._parallel(self.fake_model).state_module() + student_model_state, student_optim_state = get_state_dict(student_module, self.optimizer, options=options) + fake_model_state, fake_optim_state = get_state_dict(fake_module, self.fake_optimizer, options=options) state = { "student_model": student_model_state, "student_optimizer": student_optim_state, @@ -290,14 +306,14 @@ def _load_distributed_state(self, resume_ckpt_path): } dcp.load(state, checkpoint_id=dist_state_path) set_state_dict( - self.model.fsdp2_state_module(), + student_module, self.optimizer, model_state_dict=state["student_model"], optim_state_dict=state["student_optimizer"], options=options, ) set_state_dict( - self.fake_model.fsdp2_state_module(), + fake_module, self.fake_optimizer, model_state_dict=state["fake_model"], optim_state_dict=state["fake_optimizer"], @@ -317,8 +333,9 @@ def _load_distributed_state(self, resume_ckpt_path): ) raise RuntimeError(f"Checkpoint declares an independent fake_real role, but its distributed state is missing: {role_path}") if restored: + fake_real_module = self._parallel(self.fake_real_model).state_module() model_state, optimizer_state = get_state_dict( - self.fake_real_model.fsdp2_state_module(), + fake_real_module, self.fake_real_optimizer, options=options, ) @@ -328,7 +345,7 @@ def _load_distributed_state(self, resume_ckpt_path): } dcp.load(role_state, checkpoint_id=role_path) set_state_dict( - self.fake_real_model.fsdp2_state_module(), + fake_real_module, self.fake_real_optimizer, model_state_dict=role_state["model"], optim_state_dict=role_state["optimizer"], @@ -380,13 +397,14 @@ def save_checkpoint(self, iteration, save_total_limit): os.makedirs(save_dir, exist_ok=True) barrier() - save_student_weights = self.student_train_type == "lora" or not self.model.is_fsdp2_wrapped() + save_student_weights = self.student_train_type == "lora" or not self.parallel.is_fsdp() if save_student_weights: self._save_model_weights(self.model, save_dir, role="student") barrier() fake_save_dir = self._fake_weights_dir(save_dir) - save_fake_weights = self.fake_train_type == "lora" or not self.fake_model.is_fsdp2_wrapped() + fake_parallel = self._parallel(self.fake_model) + save_fake_weights = self.fake_train_type == "lora" or not fake_parallel.is_fsdp() if save_fake_weights and is_main_process(): os.makedirs(fake_save_dir, exist_ok=True) barrier() @@ -399,7 +417,7 @@ def save_checkpoint(self, iteration, save_total_limit): save_dir, self.role_registry.weight_directory_name(role), ) - save_fake_real_weights = self.fake_real_train_type == "lora" or not self.fake_real_model.is_fsdp2_wrapped() + save_fake_real_weights = self.fake_real_train_type == "lora" or not self._parallel(self.fake_real_model).is_fsdp() if save_fake_real_weights and is_main_process(): os.makedirs(fake_real_save_dir, exist_ok=True) barrier() @@ -421,7 +439,7 @@ def save_checkpoint(self, iteration, save_total_limit): if is_main_process() and config_path is not None: shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) - if self.model.is_fsdp2_wrapped() or self.fake_model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp() or fake_parallel.is_fsdp(): self._save_distributed_state(save_dir, iteration) if self._should_save_consolidated_student(): self._save_consolidated_student_weights(save_dir) @@ -474,7 +492,7 @@ def _should_save_consolidated_student(self): def _save_consolidated_student_weights(self, save_dir): output_dir = os.path.join(save_dir, "student_consolidated") logger.info("[train] saving consolidated student weights to {}", output_dir) - self.model.save_full_model(output_dir) + self._checkpoint(self.model).save_full_model(output_dir) barrier() def _save_distributed_state(self, save_dir, iteration): @@ -501,8 +519,16 @@ def _save_distributed_state(self, save_dir, iteration): barrier() options = StateDictOptions(ignore_frozen_params=True, strict=False) - student_model_state, student_optim_state = get_state_dict(self.model.fsdp2_state_module(), self.optimizer, options=options) - fake_model_state, fake_optim_state = get_state_dict(self.fake_model.fsdp2_state_module(), self.fake_optimizer, options=options) + student_model_state, student_optim_state = get_state_dict( + self.parallel.state_module(), + self.optimizer, + options=options, + ) + fake_model_state, fake_optim_state = get_state_dict( + self._parallel(self.fake_model).state_module(), + self.fake_optimizer, + options=options, + ) state = { "student_model": student_model_state, "student_optimizer": student_optim_state, @@ -513,7 +539,7 @@ def _save_distributed_state(self, save_dir, iteration): if getattr(self, "fake_real_model", None) is not None: role_path = os.path.join(dist_state_path, "fake_real") model_state, optimizer_state = get_state_dict( - self.fake_real_model.fsdp2_state_module(), + self._parallel(self.fake_real_model).state_module(), self.fake_real_optimizer, options=options, ) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/config.py b/lightx2v_train/lightx2v_train/trainers/dmd/config.py index 1de73efa8..53162fa86 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/config.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/config.py @@ -6,6 +6,8 @@ Optional, ) +from lightx2v_train.utils.image_size_buckets import parse_image_size_buckets + def _lora_config( role_config: Dict[str, Any], @@ -35,7 +37,7 @@ class DmdConfig: num_inference_steps: int fake_update_ratio: int guidance_scale: float - negative_prompt: str + negative_prompt: Optional[str] cfg_norm: str image_sizes: list random_schedule_enabled: bool @@ -67,14 +69,11 @@ def from_mapping( "negative_prompt", dmd.get("negative_prompt"), ) - if default_negative_prompt is not None: - negative_prompt = default_negative_prompt - elif configured_negative_prompt is not None: - negative_prompt = configured_negative_prompt - else: - raise ValueError("DMD training requires training.teacher.negative_prompt for this model.") + negative_prompt = default_negative_prompt if default_negative_prompt is not None else configured_negative_prompt random_schedule = dmd.get("random_schedule", {}) + image_sizes = dmd.get("image_sizes", []) + parse_image_size_buckets(image_sizes) return cls( student=student, fake=fake, @@ -100,7 +99,7 @@ def from_mapping( "cfg_norm", dmd.get("cfg_norm", "layer_norm"), ), - image_sizes=dmd.get("image_sizes", []), + image_sizes=image_sizes, random_schedule_enabled=bool(random_schedule.get("enabled", False)), random_schedule_num_steps_min=int(random_schedule.get("num_steps_min", 1)), random_schedule_num_steps_max=int( @@ -119,8 +118,8 @@ def from_mapping( @dataclass(frozen=True) -class VideoDmdConfig: - """Parsed video schedule and student checkpoint configuration.""" +class DmdScheduleConfig: + """Rollout schedule and optional student checkpoint configuration.""" num_train_timestep: int denoising_step_list: list diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/ltx_trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/ltx_trainer.py deleted file mode 100644 index ad4249f32..000000000 --- a/lightx2v_train/lightx2v_train/trainers/dmd/ltx_trainer.py +++ /dev/null @@ -1,675 +0,0 @@ -import math -import os - -import torch -from loguru import logger - -from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value -from lightx2v_train.utils.constants import LTX2_NEGATIVE_PROMPT -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .math import ( - detach_pair, - dmd_loss_pair, - to_dtype_pair, - velocity_to_x0, - weighted_mse_pair, -) -from .video_ar_trainer import VideoArDmdTrainer -from .video_trainer import VideoDmdTrainer - - -class _LTX2T2AVDmdMixin: - trainer_name = "ltx_t2av_dmd" - allowed_model_names = {"ltx_t2av"} - default_negative_prompt = LTX2_NEGATIVE_PROMPT - default_lora_target_modules = None - supports_diversity_loss = False - supports_real_data_fake = False - - def __init__(self, config): - from lightx2v_train.model_zoo.native.ltx2 import ( - AudioLatentShape, - AudioPatchifier, - Modality, - SpatioTemporalScaleFactors, - VideoLatentPatchifier, - VideoLatentShape, - get_pixel_coords, - ) - - self._ltx_audio_latent_shape = AudioLatentShape - self._ltx_modality = Modality - self._ltx_scale_factors = SpatioTemporalScaleFactors - self._ltx_video_latent_shape = VideoLatentShape - self._ltx_get_pixel_coords = get_pixel_coords - super().__init__(config) - model_name = self.model_config["name"] - if model_name not in self.allowed_model_names: - allowed = ", ".join(repr(name) for name in sorted(self.allowed_model_names)) - raise ValueError(f"{self.trainer_name} trainer requires model.name in {{{allowed}}}.") - - ltx2_config = self.training_config.get("ltx2", {}) - self.video_loss_weight = float(ltx2_config.get("video_loss_weight", 1.0)) - self.audio_loss_weight = float(ltx2_config.get("audio_loss_weight", 1.0)) - self.default_fps = float(ltx2_config.get("default_fps", self.dmd_config.get("fps", 25.0))) - self.video_patchifier = VideoLatentPatchifier(patch_size=1) - self.audio_patchifier = AudioPatchifier(patch_size=1) - self.denoising_sigma_values = self._validate_denoising_sigma_values(self.dmd_config.get("denoising_sigma_values")) - self._logged_denoising_sigma_values = False - - if self.cdm_enabled: - raise ValueError("ltx_t2av_dmd does not support training.dmd.cdm yet.") - - def _load_student_checkpoint(self, checkpoint_path, strict=True): - loaded_path = getattr(self.model, "loaded_student_checkpoint_path", None) - requested_path = os.path.abspath(os.path.expanduser(str(checkpoint_path))) - if loaded_path == requested_path: - logger.info("[train] LTX2 AR student checkpoint was loaded during model construction: {}", requested_path) - return - super()._load_student_checkpoint(checkpoint_path, strict=strict) - - def _encode_conditions(self, sample): - conditioning = sample["conditioning"] - positive = conditioning.get("positive") - if positive is None: - return super()._encode_conditions(sample) - - with torch.no_grad(): - condition = self._prepare_cached_condition(positive) - if self.guidance_scale > 1: - negative = conditioning.get("negative") - if negative is not None: - negative_condition = self._prepare_cached_condition(negative) - else: - prompt = conditioning.get("prompt", "") - is_scalar_prompt = isinstance(prompt, str) - batch_size = 1 if is_scalar_prompt else len(prompt) - negative_prompt = self._negative_prompt_for_conditioning( - conditioning, - batch_size, - return_scalar=is_scalar_prompt, - ) - try: - negative_condition = self.model.encode_prompt_condition(negative_prompt) - except RuntimeError as exc: - raise RuntimeError( - "LTX DMD CFG needs negative conditions. Provide data.train.negative_condition_path or keep model.load_text_encoder=true for online negative prompt encoding." - ) from exc - else: - negative_condition = None - - condition = broadcast_sequence_parallel_value(condition) - negative_condition = broadcast_sequence_parallel_value(negative_condition) if negative_condition is not None else None - return condition, negative_condition - - def _prepare_cached_condition(self, conditions): - if isinstance(conditions, (list, tuple)) and len(conditions) == 3: - video_context, audio_context, context_mask = conditions - return ( - video_context.to(device=self.model.device, dtype=self.running_dtype), - audio_context.to(device=self.model.device, dtype=self.running_dtype), - context_mask.to(device=self.model.device), - ) - - if isinstance(conditions, dict) and {"video_context", "audio_context", "context_mask"}.issubset(conditions): - return ( - conditions["video_context"].to(device=self.model.device, dtype=self.running_dtype), - conditions["audio_context"].to(device=self.model.device, dtype=self.running_dtype), - conditions["context_mask"].to(device=self.model.device), - ) - - return self.model.prepare_text_condition(conditions) - - def _latent_shape(self, sample): - prompt = sample["conditioning"].get("prompt", "") - batch_size = 1 if isinstance(prompt, str) else len(prompt) - dmd_config = self.dmd_config - - video_shape = dmd_config.get("video_latent_shape") - if video_shape is not None: - video_shape = self._shape_with_batch(video_shape, batch_size, "video_latent_shape") - else: - meta = sample["meta"] - height = int(meta["target_height"][0].item()) if "target_height" in meta else int(dmd_config.get("height", 512)) - width = int(meta["target_width"][0].item()) if "target_width" in meta else int(dmd_config.get("width", 768)) - num_frames = int(dmd_config.get("num_frames", 241)) - latent_frames = (num_frames - 1) // int(dmd_config.get("video_temporal_scale", 8)) + 1 - latent_height = int(height) // int(dmd_config.get("video_spatial_scale", 32)) - latent_width = int(width) // int(dmd_config.get("video_spatial_scale", 32)) - video_shape = ( - batch_size, - int(dmd_config.get("video_latent_channels", 128)), - latent_frames, - latent_height, - latent_width, - ) - - audio_shape = dmd_config.get("audio_latent_shape") - if audio_shape is not None: - audio_shape = self._shape_with_batch(audio_shape, batch_size, "audio_latent_shape") - else: - if "audio_num_frames" in dmd_config: - audio_frames = int(dmd_config["audio_num_frames"]) - else: - pixel_frames = int(dmd_config.get("num_frames", (video_shape[2] - 1) * 8 + 1)) - fps = float(dmd_config.get("fps", self.default_fps)) - audio_frames = round((pixel_frames / fps) * 25.0) - audio_shape = ( - batch_size, - int(dmd_config.get("audio_latent_channels", 8)), - audio_frames, - int(dmd_config.get("audio_mel_bins", 16)), - ) - - video_tokens = video_shape[2] * video_shape[3] * video_shape[4] - audio_tokens = audio_shape[2] - return { - "video_latent": video_shape, - "audio_latent": audio_shape, - "video_tokens": (batch_size, video_tokens, video_shape[1]), - "audio_tokens": (batch_size, audio_tokens, audio_shape[1] * audio_shape[3]), - "fps": float(dmd_config.get("fps", self.default_fps)), - } - - def _prepare_sampling_schedule(self, latent_shape): - if self.random_schedule_enabled: - num_steps = self._sample_synced_int(self.random_schedule_num_steps_min, self.random_schedule_num_steps_max + 1) - self.scheduler.set_random_timesteps( - self.random_schedule_num_steps_min, - self.random_schedule_num_steps_max, - sigma_min=self.random_schedule_sigma_min, - sigma_max=self.random_schedule_sigma_max, - sampling_method=self.random_schedule_sampling_method, - device=self.model.device, - num_steps=num_steps, - ) - return - sigmas = None if self.denoising_sigma_values is None else self.denoising_sigma_values[:-1] - self.scheduler.set_timesteps( - self.num_inference_steps, - sigmas=sigmas, - device=self.model.device, - ) - if not self._logged_denoising_sigma_values: - logger.info( - "[train] {} denoising_sigma_values={}", - self.trainer_name, - [round(float(value), 6) for value in self.scheduler.sigmas.detach().cpu()], - ) - self._logged_denoising_sigma_values = True - - def _validate_denoising_sigma_values(self, values): - if values is None: - return None - values = tuple(float(value) for value in values) - expected = self.num_inference_steps + 1 - if len(values) != expected: - raise ValueError(f"training.dmd.denoising_sigma_values must contain num_inference_steps + 1 values ({expected}), got {len(values)}.") - if abs(values[0] - 1.0) > 1e-6 or abs(values[-1]) > 1e-6: - raise ValueError("training.dmd.denoising_sigma_values must start at 1.0 and end at 0.0.") - if any(current <= following for current, following in zip(values, values[1:])): - raise ValueError("training.dmd.denoising_sigma_values must be strictly decreasing.") - return values - - def sample_initial_latents(self, latent_shape): - video = torch.randn(latent_shape["video_tokens"], device=self.model.device, dtype=self.running_dtype) - audio = torch.randn(latent_shape["audio_tokens"], device=self.model.device, dtype=self.running_dtype) - return video, audio - - def run_back_simulation(self, condition, latent_shape, end_step_idx, grad_enabled, xt=None): - if xt is None: - xt = self.sample_initial_latents(latent_shape) - x0 = None - xt_end = None - vt_end = None - self.model.transformer.train() - for idx in range(end_step_idx + 1): - sigma = self.scheduler.sigma_at(idx, latent_shape["video_tokens"][0], device=self.model.device, dtype=self.running_dtype) - context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad - with context(): - velocity = self._predict_velocity(self.model, xt, sigma, condition, latent_shape) - if idx == end_step_idx: - xt_end = self._detach_pair(xt) - vt_end = self._detach_pair(velocity) - xt, x0 = self._step_pair_by_index(velocity, idx, xt) - return self._to_dtype_pair(x0, self.running_dtype), xt_end, vt_end - - def forward_loss( - self, - latent_shape, - conditions, - stage, - current_iter=None, - initial_noise=None, - ): - condition, negative_condition = conditions - self._prepare_sampling_schedule(latent_shape) - end_step_idx = self.sample_end_step() - xt_start = self.sample_initial_latents(latent_shape) if initial_noise is None else initial_noise - x0, xt_end, vt_end = self.run_back_simulation(condition, latent_shape, end_step_idx, grad_enabled=(stage != "fake"), xt=xt_start) - - batch_size = latent_shape["video_tokens"][0] - sigma = self.scheduler.sample_renoise_sigma(batch_size, device=self.model.device, dtype=self.running_dtype) - noise = ( - torch.randn_like(x0[0], dtype=torch.float32), - torch.randn_like(x0[1], dtype=torch.float32), - ) - renoised_xt = self._add_noise_pair(x0, noise, sigma) - - if stage == "fake": - self.fake_model.transformer.train() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition, latent_shape) - velocity_gt = (noise[0] - x0[0].float(), noise[1] - x0[1].float()) - return self._weighted_mse_pair(velocity_fake, velocity_gt) - - with torch.no_grad(): - self.fake_model.transformer.eval() - self.teacher_model.transformer.eval() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition, latent_shape) - velocity_teacher_cond = self._predict_velocity(self.teacher_model, renoised_xt, sigma, condition, latent_shape) - if negative_condition is None or self.guidance_scale == 0: - velocity_teacher = velocity_teacher_cond - else: - velocity_teacher_uncond = self._predict_velocity(self.teacher_model, renoised_xt, sigma, negative_condition, latent_shape) - velocity_teacher = self._cfg_pair(velocity_teacher_cond, velocity_teacher_uncond) - - x_pred_fake = self._x0_from_velocity(renoised_xt, velocity_fake, sigma) - x_pred_teacher = self._x0_from_velocity(renoised_xt, velocity_teacher, sigma) - return self._dmd_loss_pair(x0, x_pred_fake, x_pred_teacher) - - def _predict_velocity(self, model, latents, sigma, condition, latent_shape): - video_context, audio_context, context_mask = condition - video_tokens, audio_tokens = latents - batch_size = video_tokens.shape[0] - sigma = sigma.to(device=self.model.device, dtype=self.running_dtype) - timesteps_video = sigma.view(-1, 1).expand(batch_size, video_tokens.shape[1]).clone() - timesteps_audio = sigma.view(-1, 1).expand(batch_size, audio_tokens.shape[1]).clone() - video_positions = self._get_video_positions(latent_shape, batch_size, video_tokens.device) - audio_positions = self._get_audio_positions(latent_shape, batch_size, audio_tokens.device) - - video_modality = self._ltx_modality( - enabled=True, - latent=video_tokens.to(device=model.device, dtype=self.running_dtype), - sigma=sigma, - timesteps=timesteps_video, - positions=video_positions, - context=video_context.to(device=model.device, dtype=self.running_dtype), - context_mask=context_mask.to(device=model.device), - ) - audio_modality = self._ltx_modality( - enabled=True, - latent=audio_tokens.to(device=model.device, dtype=self.running_dtype), - sigma=sigma, - timesteps=timesteps_audio, - positions=audio_positions, - context=audio_context.to(device=model.device, dtype=self.running_dtype), - context_mask=context_mask.to(device=model.device), - ) - with model.transformer_forward_context(): - return model.denoiser_module()( - video=video_modality, - audio=audio_modality, - perturbations=None, - ) - - def _get_video_positions(self, latent_shape, batch_size, device): - _, channels, frames, height, width = latent_shape["video_latent"] - latent_coords = self.video_patchifier.get_patch_grid_bounds( - output_shape=self._ltx_video_latent_shape( - frames=frames, - height=height, - width=width, - batch=batch_size, - channels=channels, - ), - device=device, - ) - pixel_coords = self._ltx_get_pixel_coords( - latent_coords=latent_coords, - scale_factors=self._ltx_scale_factors.default(), - causal_fix=True, - ).float() - pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / latent_shape["fps"] - return pixel_coords - - def _get_audio_positions(self, latent_shape, batch_size, device): - _, channels, frames, mel_bins = latent_shape["audio_latent"] - return self.audio_patchifier.get_patch_grid_bounds( - output_shape=self._ltx_audio_latent_shape( - frames=frames, - mel_bins=mel_bins, - batch=batch_size, - channels=channels, - ), - device=device, - ) - - def _add_noise_pair(self, latents, noises, sigma): - return ( - self.scheduler.add_noise(latents[0], noises[0], sigma).to(dtype=latents[0].dtype), - self.scheduler.add_noise(latents[1], noises[1], sigma).to(dtype=latents[1].dtype), - ) - - def _step_pair_by_index(self, velocity, step_idx, sample): - video_next, video_x0 = self.scheduler.step_by_index(velocity[0], step_idx, sample[0]) - audio_next, audio_x0 = self.scheduler.step_by_index(velocity[1], step_idx, sample[1]) - return (video_next, audio_next), (video_x0, audio_x0) - - def _x0_from_velocity(self, xt, velocity, sigma): - return ( - velocity_to_x0(xt[0], velocity[0], sigma), - velocity_to_x0(xt[1], velocity[1], sigma), - ) - - def _cfg_pair(self, cond, uncond): - return ( - self._do_cfg(cond[0], uncond[0], self.guidance_scale, self.cfg_norm), - self._do_cfg(cond[1], uncond[1], self.guidance_scale, self.cfg_norm), - ) - - def _dmd_loss_pair(self, latents, x_pred_fake_flow, x_pred_teacher): - return dmd_loss_pair( - latents, - x_pred_fake_flow, - x_pred_teacher, - self.video_loss_weight, - self.audio_loss_weight, - ) - - def _weighted_mse_pair(self, pred, target): - return weighted_mse_pair( - pred, - target, - self.video_loss_weight, - self.audio_loss_weight, - ) - - @staticmethod - def _detach_pair(value): - return detach_pair(value) - - @staticmethod - def _to_dtype_pair(value, dtype): - return to_dtype_pair(value, dtype) - - @staticmethod - def _shape_with_batch(shape, batch_size, name): - shape = tuple(int(dim) for dim in shape) - if name.startswith("audio") and len(shape) == 3: - return (batch_size, *shape) - if name.startswith("audio") and len(shape) == 4: - if shape[0] != batch_size: - return (batch_size, *shape[1:]) - return shape - if name.startswith("video") and len(shape) == 4: - return (batch_size, *shape) - if name.startswith("video") and len(shape) == 5: - if shape[0] != batch_size: - return (batch_size, *shape[1:]) - return shape - raise ValueError(f"training.dmd.{name} has unsupported shape {shape}.") - - -@TRAINER_REGISTER("ltx_t2av_dmd") -class LTX2T2AVDmdTrainer(_LTX2T2AVDmdMixin, VideoDmdTrainer): - trainer_name = "ltx_t2av_dmd" - - -@TRAINER_REGISTER("ltx_t2av_ar_dmd") -class LTX2T2AVArDmdTrainer(_LTX2T2AVDmdMixin, VideoArDmdTrainer): - trainer_name = "ltx_t2av_ar_dmd" - allowed_model_names = {"ltx_t2av_ar"} - - def __init__(self, config): - super().__init__(config) - if self.denoising_sigma_values is None: - raise ValueError("ltx_t2av_ar_dmd requires training.dmd.denoising_sigma_values so training and inference use the same distilled trajectory.") - - def run_back_simulation(self, condition, latent_shape, end_step_idx, grad_enabled, xt=None): - transformer = self.model.denoiser_module() - if not self._ltx_transformer_has_inference_cache(transformer): - raise RuntimeError("ltx_t2av_ar_dmd requires an LTX transformer with _forward_inference/KV cache support.") - - if xt is None: - xt = self.sample_initial_latents(latent_shape) - video_tokens, audio_tokens = xt - batch_size = video_tokens.shape[0] - _, _, num_frames, latent_height, latent_width = latent_shape["video_latent"] - frame_seq_length = int(latent_height) * int(latent_width) - if video_tokens.shape[1] != num_frames * frame_seq_length: - if video_tokens.shape[1] % num_frames != 0: - raise ValueError(f"LTX AR-DMD video token length {video_tokens.shape[1]} is not divisible by latent frames {num_frames}.") - frame_seq_length = video_tokens.shape[1] // num_frames - if num_frames % self.num_frame_per_chunk != 0: - raise ValueError(f"LTX AR-DMD latent frames={num_frames} must be divisible by num_frame_per_chunk={self.num_frame_per_chunk}.") - - self.model.transformer.train() - num_blocks = num_frames // self.num_frame_per_chunk - video_positions = self._get_video_positions(latent_shape, batch_size, video_tokens.device) - audio_positions = self._get_audio_positions(latent_shape, batch_size, audio_tokens.device) - kv_cache = self._new_ltx_kv_caches( - batch_size=batch_size, - dtype=video_tokens.dtype, - device=video_tokens.device, - video_total_tokens=video_tokens.shape[1], - audio_total_tokens=audio_tokens.shape[1], - video_frame_seq_length=frame_seq_length, - num_frames=num_frames, - ) - - output_video_chunks = [] - output_audio_chunks = [] - for block_idx in range(num_blocks): - frame_start = block_idx * self.num_frame_per_chunk - frame_end = frame_start + self.num_frame_per_chunk - video_start = frame_start * frame_seq_length - video_end = frame_end * frame_seq_length - audio_start, audio_end = self._audio_token_range(block_idx, num_blocks, audio_tokens.shape[1]) - latents = ( - video_tokens[:, video_start:video_end], - audio_tokens[:, audio_start:audio_end], - ) - pos = ( - video_positions[:, :, video_start:video_end], - audio_positions[:, :, audio_start:audio_end], - ) - - x0 = None - for idx in range(end_step_idx + 1): - sigma = self.scheduler.sigma_at(idx, batch_size, device=self.model.device, dtype=self.running_dtype) - context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad - with context(): - velocity = self._predict_velocity_chunk( - self.model, - latents, - sigma, - condition, - pos, - kv_cache=kv_cache, - video_current_start=video_start, - audio_current_start=audio_start, - ) - x0 = self._x0_from_velocity(latents, velocity, sigma) - - if idx < end_step_idx: - next_sigma = self.scheduler.sigma_at( - idx + 1, - batch_size, - device=self.model.device, - dtype=self.running_dtype, - ) - with torch.no_grad(): - noise = ( - torch.randn_like(x0[0]), - torch.randn_like(x0[1]), - ) - latents = self._add_noise_pair( - x0, - noise, - next_sigma, - ) - - output_video_chunks.append(x0[0]) - output_audio_chunks.append(x0[1]) - - cache_latents = self._detach_pair(x0) - if self.context_noise > 0: - cache_sigma = torch.full((batch_size,), self.context_noise, device=self.model.device, dtype=self.running_dtype) - cache_noise = (torch.randn_like(cache_latents[0]), torch.randn_like(cache_latents[1])) - cache_latents = self._add_noise_pair(cache_latents, cache_noise, cache_sigma) - else: - cache_sigma = torch.zeros((batch_size,), device=self.model.device, dtype=self.running_dtype) - with torch.no_grad(): - self._predict_velocity_chunk( - self.model, - cache_latents, - cache_sigma, - condition, - pos, - kv_cache=kv_cache, - video_current_start=video_start, - audio_current_start=audio_start, - ) - - return ( - self._to_dtype_pair((torch.cat(output_video_chunks, dim=1), torch.cat(output_audio_chunks, dim=1)), self.running_dtype), - None, - None, - ) - - def _predict_velocity_chunk( - self, - model, - latents, - sigma, - condition, - positions, - *, - kv_cache=None, - video_current_start=0, - audio_current_start=0, - ): - video_context, audio_context, context_mask = condition - video_tokens, audio_tokens = latents - video_positions, audio_positions = positions - batch_size = video_tokens.shape[0] - sigma = sigma.to(device=model.device, dtype=self.running_dtype) - video_modality = self._ltx_modality( - enabled=True, - latent=video_tokens.to(device=model.device, dtype=self.running_dtype), - sigma=sigma, - timesteps=sigma.view(-1, 1).expand(batch_size, video_tokens.shape[1]).clone(), - positions=video_positions.to(device=model.device), - context=video_context.to(device=model.device, dtype=self.running_dtype), - context_mask=context_mask.to(device=model.device), - ) - audio_modality = self._ltx_modality( - enabled=True, - latent=audio_tokens.to(device=model.device, dtype=self.running_dtype), - sigma=sigma, - timesteps=sigma.view(-1, 1).expand(batch_size, audio_tokens.shape[1]).clone(), - positions=audio_positions.to(device=model.device), - context=audio_context.to(device=model.device, dtype=self.running_dtype), - context_mask=context_mask.to(device=model.device), - ) - with model.transformer_forward_context(): - return model.denoiser_module()( - video=video_modality, - audio=audio_modality, - perturbations=None, - kv_cache=kv_cache, - video_current_start=video_current_start, - audio_current_start=audio_current_start, - ) - - def _new_ltx_kv_caches( - self, - batch_size, - dtype, - device, - video_total_tokens, - audio_total_tokens, - video_frame_seq_length, - num_frames, - ): - transformer = self._unwrap_ltx_transformer(self.model.denoiser_module()) - num_layers = len(transformer.transformer_blocks) - video_heads = int(transformer.num_attention_heads) - video_head_dim = int(transformer.inner_dim // transformer.num_attention_heads) - audio_heads = int(transformer.audio_num_attention_heads) - audio_head_dim = int(transformer.audio_inner_dim // transformer.audio_num_attention_heads) - local_attn_size = int(getattr(transformer, "local_attn_size", -1)) - sink_size = int(getattr(transformer, "sink_size", 0)) - video_chunk_tokens = self.num_frame_per_chunk * video_frame_seq_length - if local_attn_size == -1: - video_cache_size = int(video_total_tokens) - video_attention_window = int(video_total_tokens) - else: - video_attention_window = max(video_chunk_tokens, local_attn_size * video_frame_seq_length) - video_cache_size = max(video_attention_window, sink_size * video_frame_seq_length + video_chunk_tokens) - - audio_tokens_per_frame = max(1, math.ceil(audio_total_tokens / max(1, num_frames))) - audio_chunk_tokens = max(1, math.ceil(audio_total_tokens / (num_frames // self.num_frame_per_chunk))) - if local_attn_size == -1: - audio_cache_size = int(audio_total_tokens) - audio_attention_window = int(audio_total_tokens) - else: - audio_attention_window = max(audio_chunk_tokens, local_attn_size * audio_tokens_per_frame) - audio_cache_size = max(audio_attention_window, sink_size * audio_tokens_per_frame + audio_chunk_tokens) - - cache = {"video": [], "audio": []} - for _ in range(num_layers): - cache["video"].append( - self._empty_ltx_cache( - batch_size, - video_cache_size, - video_heads, - video_head_dim, - dtype, - device, - attention_window_size=video_attention_window, - sink_tokens=sink_size * video_frame_seq_length, - ) - ) - cache["audio"].append( - self._empty_ltx_cache( - batch_size, - audio_cache_size, - audio_heads, - audio_head_dim, - dtype, - device, - attention_window_size=audio_attention_window, - sink_tokens=sink_size * audio_tokens_per_frame, - ) - ) - return cache - - @staticmethod - def _empty_ltx_cache(batch_size, cache_size, heads, head_dim, dtype, device, *, attention_window_size, sink_tokens): - return { - "k": torch.zeros((batch_size, cache_size, heads, head_dim), dtype=dtype, device=device), - "v": torch.zeros((batch_size, cache_size, heads, head_dim), dtype=dtype, device=device), - "global_end_index": torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": torch.tensor([0], dtype=torch.long, device=device), - "attention_window_size": int(attention_window_size), - "sink_tokens": int(sink_tokens), - } - - @staticmethod - def _audio_token_range(block_idx, num_blocks, audio_total_tokens): - start = (block_idx * audio_total_tokens + num_blocks - 1) // num_blocks - end = ((block_idx + 1) * audio_total_tokens + num_blocks - 1) // num_blocks - return start, end - - @staticmethod - def _ltx_transformer_has_inference_cache(transformer): - return LTX2T2AVArDmdTrainer._unwrap_ltx_transformer(transformer) is not None - - @staticmethod - def _unwrap_ltx_transformer(transformer): - for candidate in (transformer, getattr(transformer, "module", None), getattr(transformer, "_fsdp_wrapped_module", None)): - if candidate is not None and hasattr(candidate, "_forward_inference"): - return candidate - return None diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py new file mode 100644 index 000000000..b849fb42e --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py @@ -0,0 +1,459 @@ +import copy + +import torch +import torch.distributed as dist +from loguru import logger + +from lightx2v_train.model_capabilities import ( + CheckpointCapability, + DistillationCapability, + ParallelCapability, + TrainableModelCapability, +) +from lightx2v_train.model_zoo import build_loaded_model +from lightx2v_train.runtime.distributed import is_distributed +from lightx2v_train.runtime.sequence_parallel import ( + broadcast_sequence_parallel_value, + sync_sequence_parallel_gradients, +) +from lightx2v_train.schedulers import DMDFlowMatchingScheduler +from lightx2v_train.tricks import ( + CdmStepContext, + CdmTrainerConstraints, + CdmTrick, + IdaModelPair, + IdaSetupContext, + IdaStepContext, + ImplicitDistributionAlignmentTrick, +) + +from ..base import BaseTrainer +from .checkpoint import DmdCheckpointManager +from .config import DmdConfig +from .math import ( + dmd_loss, + do_cfg, +) +from .roles import DmdRoleRegistry + + +class _DmdRuntime(BaseTrainer): + trainer_name = "dmd" + required_capabilities = ( + *BaseTrainer.required_capabilities, + DistillationCapability, + ) + default_negative_prompt = None + supports_cdm = True + supports_ida = True + defer_ida_setup = False + + def _resolve_train_type(self): + if "train_type" in self.training_config: + raise ValueError("DMD trainers use training.student.train_type and training.fake.train_type; remove training.train_type.") + return None + + def __init__(self, config): + super().__init__(config) + self.role_registry = DmdRoleRegistry(self) + self.checkpoint_manager = DmdCheckpointManager(self) + parsed = DmdConfig.from_mapping( + config, + default_negative_prompt=self.default_negative_prompt, + ) + self.parsed_dmd_config = parsed + self.student_config = parsed.student + self.fake_config = parsed.fake + self.student_train_type = parsed.student_train_type + self.fake_train_type = parsed.fake_train_type + self.student_lora_config = parsed.student_lora + self.fake_lora_config = parsed.fake_lora + + self.fake_optimizer_config = self.fake_config["optimizer"] + self.fake_optimizer_learning_rate = self.fake_optimizer_config.get("learning_rate", self.optimizer_learning_rate) + self.fake_optimizer_adam_beta1 = self.fake_optimizer_config.get("adam_beta1", self.optimizer_adam_beta1) + self.fake_optimizer_adam_beta2 = self.fake_optimizer_config.get("adam_beta2", self.optimizer_adam_beta2) + self.fake_optimizer_weight_decay = self.fake_optimizer_config.get("weight_decay", self.optimizer_weight_decay) + self.fake_optimizer_adam_epsilon = self.fake_optimizer_config.get("adam_epsilon", self.optimizer_adam_epsilon) + + self.dmd_config = parsed.dmd + self.cdm_trick = CdmTrick.from_mapping(self.dmd_config.get("cdm", {})) + self.cdm_trick.validate_trainer( + CdmTrainerConstraints( + supported=self.supports_cdm, + ), + self.trainer_name, + ) + self.ida_trick = ImplicitDistributionAlignmentTrick.from_mapping(self.dmd_config.get("ida", {})) + if self.ida_trick.enabled and not self.supports_ida: + raise ValueError(f"{self.trainer_name} does not support training.dmd.ida.") + self.num_inference_steps = parsed.num_inference_steps + self.fake_update_ratio = parsed.fake_update_ratio + self.guidance_scale = parsed.guidance_scale + self.negative_prompt = parsed.negative_prompt + self.cfg_norm = parsed.cfg_norm + self.image_sizes = parsed.image_sizes + self.random_schedule_enabled = parsed.random_schedule_enabled + self.random_schedule_num_steps_min = parsed.random_schedule_num_steps_min + self.random_schedule_num_steps_max = parsed.random_schedule_num_steps_max + self.random_schedule_sigma_min = parsed.random_schedule_sigma_min + self.random_schedule_sigma_max = parsed.random_schedule_sigma_max + self.random_schedule_sampling_method = parsed.random_schedule_sampling_method + + def set_model(self, model): + super().set_model(model) + self.student = model.capabilities.require(DistillationCapability) + if self.negative_prompt is None: + self.negative_prompt = self.student.default_negative_prompt + for lora_config in ( + self.student_lora_config, + self.fake_lora_config, + ): + if lora_config is not None and "target_modules" not in lora_config and self.student.default_lora_target_modules is not None: + lora_config["target_modules"] = list(self.student.default_lora_target_modules) + + @property + def cdm_enabled(self): + return self.cdm_trick.enabled + + def _get_optimizer_config(self): + return self.training_config["student"]["optimizer"] + + def _setup_trainable_model(self, model, role="student"): + if role == "student": + train_type = self.student_train_type + lora_config = self.student_lora_config + elif role in {"fake", "fake_real"}: + train_type = self.fake_train_type + lora_config = self.fake_lora_config + else: + raise ValueError(f"Unsupported DMD model role: {role}") + model.ensure_capabilities().require(TrainableModelCapability).configure(train_type, lora_config or {}) + + def _restore_trainable_model(self, model, role="student"): + if role == "student": + train_type = self.student_train_type + elif role in {"fake", "fake_real"}: + train_type = self.fake_train_type + else: + raise ValueError(f"Unsupported DMD model role: {role}") + model.ensure_capabilities().require(TrainableModelCapability).restore(train_type) + + def _save_model_weights(self, model, save_dir, role="student"): + train_type = self.student_train_type if role == "student" else self.fake_train_type + model.ensure_capabilities().require(CheckpointCapability).save_weights(save_dir, train_type) + + def _load_model_weights(self, model, save_dir, role="student"): + train_type = self.student_train_type if role == "student" else self.fake_train_type + model.ensure_capabilities().require(CheckpointCapability).load_weights(save_dir, train_type) + + def setup(self, resume_ckpt_path=None): + super().setup(resume_ckpt_path=None) + base_model_config = { + key: copy.deepcopy(value) + for key, value in self.model_config.items() + if key + not in { + "fake", + "teacher", + "student", + "student_2", + "fake_2", + "fake_low_high", + "fake_real", + "fake_real_high", + "fake_real_low", + "teacher_2", + } + } + + fake_model_config = copy.deepcopy(self.config) + fake_model_config["model"] = copy.deepcopy(base_model_config) + if "fake" in self.model_config: + if not isinstance(self.model_config["fake"], dict): + raise ValueError("model.fake must be a mapping.") + fake_model_config["model"].update(copy.deepcopy(self.model_config["fake"])) + self.fake_model_config = copy.deepcopy(fake_model_config) + self.fake_model = build_loaded_model( + fake_model_config, + transformer_only=True, + reference_model=self.model, + ) + self.fake = self.fake_model.capabilities.require(DistillationCapability) + self._setup_trainable_model(self.fake_model, role="fake") + self.fake_model.capabilities.require(ParallelCapability).apply(self.config) + if self.gradient_checkpointing: + self.fake_model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() + + teacher_model_config = copy.deepcopy(self.config) + teacher_model_config["model"] = copy.deepcopy(base_model_config) + if "teacher" in self.model_config: + if not isinstance(self.model_config["teacher"], dict): + raise ValueError("model.teacher must be a mapping.") + teacher_model_config["model"].update(copy.deepcopy(self.model_config["teacher"])) + self.teacher_model = build_loaded_model( + teacher_model_config, + transformer_only=True, + reference_model=self.model, + ) + self.teacher = self.teacher_model.capabilities.require(DistillationCapability) + self.teacher.denoiser().requires_grad_(False) + self.teacher.set_training(False) + self.teacher_model.capabilities.require(ParallelCapability).apply(self.config) + self.teacher.set_training(False) + + self.fake_trainable_params = list(self.fake_model.capabilities.require(TrainableModelCapability).parameters()) + self.fake_optimizer = self._build_optimizer( + self.fake_trainable_params, + { + "learning_rate": self.fake_optimizer_learning_rate, + "adam_beta1": self.fake_optimizer_adam_beta1, + "adam_beta2": self.fake_optimizer_adam_beta2, + "weight_decay": self.fake_optimizer_weight_decay, + "adam_epsilon": self.fake_optimizer_adam_epsilon, + }, + ) + self.fake_lr_scheduler = self._build_lr_scheduler( + self.fake_optimizer, + num_warmup_steps=0, + num_training_steps=max(1, self.max_train_iters * self.fake_update_ratio), + ) + + self.scheduler = DMDFlowMatchingScheduler(self.config, self.dmd_config) + + if resume_ckpt_path is not None: + self._load_resume_state(resume_ckpt_path) + + if not self.defer_ida_setup: + self._setup_ida_trick() + + logger.info("[train] dmd student model={} path={}", self.model_config["name"], self.model_config["pretrained_model_name_or_path"]) + logger.info("[train] dmd fake model={} path={}", fake_model_config["model"]["name"], fake_model_config["model"]["pretrained_model_name_or_path"]) + logger.info("[train] dmd teacher model={} path={}", teacher_model_config["model"]["name"], teacher_model_config["model"]["pretrained_model_name_or_path"]) + logger.info("[train] dmd train_types student={} fake={}", self.student_train_type, self.fake_train_type) + logger.info("[train] dmd student trainable params={}", self._count_trainable(self.student.denoiser())) + logger.info("[train] dmd fake trainable params={}", self._count_trainable(self.fake.denoiser())) + if self.random_schedule_enabled: + logger.info( + "[train] dmd random sigma schedule enabled: steps=[{}, {}], sigma=[{}, {}], sampling_method={}", + self.random_schedule_num_steps_min, + self.random_schedule_num_steps_max, + self.random_schedule_sigma_min, + self.random_schedule_sigma_max, + self.random_schedule_sampling_method, + ) + if self.cdm_enabled: + logger.info( + "[train] dmd CDM enabled: weight={} warmup_iters={}", + self.cdm_trick.config.weight, + self.cdm_trick.config.warmup_iters, + ) + + @staticmethod + def _count_trainable(module): + return sum(1 for param in module.parameters() if param.requires_grad) + + def _ida_model_pairs(self): + return { + "main": IdaModelPair( + student=self.student.denoiser(), + fake=self.fake.denoiser(), + ) + } + + def _setup_ida_trick(self): + self.ida_trick.setup(IdaSetupContext(model_pairs=self._ida_model_pairs())) + logger.info( + "[train] {} SenseFlow IDA enabled={} decay={}", + self.trainer_name, + self.ida_trick.enabled, + self.ida_trick.config.decay, + ) + + def _after_student_optimizer_step(self, role): + self.ida_trick.after_student_step(IdaStepContext(role=role)) + + @staticmethod + def _do_cfg(cond_pred, uncond_pred, cfg_scale, cfg_norm): + return do_cfg( + cond_pred, + uncond_pred, + cfg_scale, + cfg_norm, + ) + + @staticmethod + def _dmd_loss(latents, x_pred_fake_flow, x_pred_teacher, norm_clip_min=None): + return dmd_loss( + latents, + x_pred_fake_flow, + x_pred_teacher, + norm_clip_min=norm_clip_min, + ) + + def _prepare_sampling_schedule(self, latent_shape): + latent_hw = self.student.latent_hw(latent_shape) + if self.random_schedule_enabled: + num_steps = self._sample_synced_int(self.random_schedule_num_steps_min, self.random_schedule_num_steps_max + 1) + self.scheduler.set_random_timesteps( + self.random_schedule_num_steps_min, + self.random_schedule_num_steps_max, + sigma_min=self.random_schedule_sigma_min, + sigma_max=self.random_schedule_sigma_max, + sampling_method=self.random_schedule_sampling_method, + latent_hw=latent_hw, + device=self.student.device, + num_steps=num_steps, + ) + return + sigma_values = self.dmd_config.get("denoising_sigma_values") + if sigma_values is not None: + sigma_values = tuple(float(value) for value in sigma_values) + expected = self.num_inference_steps + 1 + if len(sigma_values) != expected: + raise ValueError(f"training.dmd.denoising_sigma_values must contain {expected} values, got {len(sigma_values)}.") + sigma_values = sigma_values[:-1] + self.scheduler.set_timesteps( + self.num_inference_steps, + sigmas=sigma_values, + latent_hw=latent_hw, + device=self.student.device, + ) + + def _latent_shape(self, sample): + return self.student.latent_shape( + sample, + self.dmd_config, + self.image_sizes, + broadcast_sequence_parallel_value, + ) + + def _encode_conditions(self, sample): + return self.student.encode_conditions( + sample, + self.negative_prompt, + self.guidance_scale, + broadcast_sequence_parallel_value, + ) + + @staticmethod + def _predict_velocity(capability, latents, sigma, condition): + return capability.predict_velocity(latents, sigma, condition) + + def _predict_teacher_velocity(self, latents, sigma, condition, negative_condition): + return self.teacher.predict_guided_velocity( + latents, + sigma, + condition, + negative_condition, + self.guidance_scale, + self.cfg_norm, + ) + + def sample_initial_latents(self, latent_shape): + return self.student.initial_latents( + latent_shape, + self.running_dtype, + broadcast_sequence_parallel_value, + ) + + def _sample_synced_int(self, low, high): + value = torch.randint(int(low), int(high), (1,), device=self.student.device, dtype=torch.int64) + if is_distributed(): + dist.broadcast(value, src=0) + return int(value.item()) + + def _cdm_step_context( + self, + xt, + vt, + end_step_idx, + condition, + current_iter, + ): + return CdmStepContext( + trajectory_latent=xt, + trajectory_velocity=vt, + end_step_index=end_step_idx, + condition=condition, + current_iteration=current_iter, + device=self.student.device, + dtype=self.running_dtype, + scheduler=self.scheduler, + predict_student_velocity=lambda latent, sigma, cond: self._predict_velocity( + self.student, + latent, + sigma, + cond, + ), + predict_fake_velocity=lambda latent, sigma, cond: self._predict_velocity( + self.fake, + latent, + sigma, + cond, + ), + predict_teacher_velocity=lambda latent, sigma, cond: self._predict_teacher_velocity( + latent, + sigma, + cond, + None, + ), + prepare_fake_model=lambda: self.fake.set_training(False), + dmd_loss=self._dmd_loss, + ) + + def _iter_train_samples(self): + epoch = 0 + while True: + sampler = getattr(self.dataloader_train, "sampler", None) + if hasattr(sampler, "set_epoch"): + sampler.set_epoch(epoch) + for sample in self.dataloader_train: + yield sample + epoch += 1 + + def _set_student_gradient_sync(self, enabled): + self.model.capabilities.require(ParallelCapability).set_gradient_sync(enabled) + + def _set_fake_gradient_sync(self, enabled): + self.fake_model.capabilities.require(ParallelCapability).set_gradient_sync(enabled) + + def _set_gradient_sync(self, enabled): + self._set_student_gradient_sync(enabled) + self._set_fake_gradient_sync(enabled) + + def _sync_sequence_parallel_grads(self, params): + sync_sequence_parallel_gradients(params) + + def _fake_weights_dir(self, root_dir): + return self.checkpoint_manager._fake_weights_dir(root_dir) + + def _trick_checkpoint_metadata(self): + return self.checkpoint_manager._trick_checkpoint_metadata() + + def _validate_optional_trick_metadata(self, state, state_path): + return self.checkpoint_manager._validate_optional_trick_metadata(state, state_path) + + def _load_resume_state(self, resume_ckpt_path): + return self.checkpoint_manager._load_resume_state(resume_ckpt_path) + + def _validate_dmd_checkpoint_metadata(self, state, state_path, resume_ckpt_path): + return self.checkpoint_manager._validate_dmd_checkpoint_metadata(state, state_path, resume_ckpt_path) + + def _load_single_process_state(self, resume_ckpt_path): + return self.checkpoint_manager._load_single_process_state(resume_ckpt_path) + + def _load_distributed_state(self, resume_ckpt_path): + return self.checkpoint_manager._load_distributed_state(resume_ckpt_path) + + def save_checkpoint(self, iteration, save_total_limit): + return self.checkpoint_manager.save_checkpoint(iteration, save_total_limit) + + def _should_save_consolidated_student(self): + return self.checkpoint_manager._should_save_consolidated_student() + + def _save_consolidated_student_weights(self, save_dir): + return self.checkpoint_manager._save_consolidated_student_weights(save_dir) + + def _save_distributed_state(self, save_dir, iteration): + return self.checkpoint_manager._save_distributed_state(save_dir, iteration) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py index 7600846f5..69ff60b7a 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py @@ -1,525 +1,250 @@ import copy import os +from functools import partial import torch -import torch.distributed as dist -import torch.nn.functional as F +import torch.distributed.checkpoint as dcp from loguru import logger +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + set_model_state_dict, +) -from lightx2v_train.model_zoo import build_model +from lightx2v_train.model_capabilities import ( + DistillationCapability, + ParallelCapability, + TrainableModelCapability, +) +from lightx2v_train.model_zoo import build_loaded_model from lightx2v_train.runtime.distributed import ( barrier, get_world_size, - is_distributed, is_main_process, reduce_mean, ) -from lightx2v_train.runtime.parallel import ( - apply_parallel, - set_parallel_gradient_sync, -) -from lightx2v_train.runtime.sequence_parallel import ( - broadcast_sequence_parallel_value, - sync_sequence_parallel_gradients, -) +from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value from lightx2v_train.schedulers import DMDFlowMatchingScheduler +from lightx2v_train.schedulers.flow_matching import CausalForcingFlowMatchScheduler from lightx2v_train.tricks import ( - CdmStepContext, - CdmTrainerConstraints, - CdmTrick, - IdaModelPair, - IdaSetupContext, - IdaStepContext, - ImplicitDistributionAlignmentTrick, + DiversitySetupContext, + DiversityStepContext, + DiversityTrainerConstraints, + DiversityTrick, + RealDataFakeSetupContext, + RealDataFakeStepContext, + RealDataFakeTrick, ) from lightx2v_train.utils.registry import TRAINER_REGISTER -from ..base import BaseTrainer -from .checkpoint import DmdCheckpointManager -from .config import DmdConfig -from .math import ( - dmd_loss, - do_cfg, -) -from .roles import DmdRoleRegistry +from .config import DmdScheduleConfig +from .runtime import _DmdRuntime @TRAINER_REGISTER("dmd") -class DmdTrainer(BaseTrainer): - trainer_name = "dmd" - default_negative_prompt = None - supports_cdm = True - supported_cdm_model_names = frozenset( - { - "flux2_dev", - "longcat_image", - "qwen_image", - } - ) - supports_ida = True - defer_ida_setup = False +class DmdTrainer(_DmdRuntime): + """Distribution matching for image, video, and joint latent models.""" - def _resolve_train_type(self): - if "train_type" in self.training_config: - raise ValueError("DMD trainers use training.student.train_type and training.fake.train_type; remove training.train_type.") - return None + trainer_name = "dmd" + supports_diversity_loss = True + supports_real_data_fake = True def __init__(self, config): super().__init__(config) - self.role_registry = DmdRoleRegistry(self) - self.checkpoint_manager = DmdCheckpointManager(self) - parsed = DmdConfig.from_mapping( + parsed = DmdScheduleConfig.from_mapping( config, - default_negative_prompt=self.default_negative_prompt, - ) - self.parsed_dmd_config = parsed - self.student_config = parsed.student - self.fake_config = parsed.fake - self.student_train_type = parsed.student_train_type - self.fake_train_type = parsed.fake_train_type - self.student_lora_config = parsed.student_lora - self.fake_lora_config = parsed.fake_lora - - self.fake_optimizer_config = self.fake_config["optimizer"] - self.fake_optimizer_learning_rate = self.fake_optimizer_config.get("learning_rate", self.optimizer_learning_rate) - self.fake_optimizer_adam_beta1 = self.fake_optimizer_config.get("adam_beta1", self.optimizer_adam_beta1) - self.fake_optimizer_adam_beta2 = self.fake_optimizer_config.get("adam_beta2", self.optimizer_adam_beta2) - self.fake_optimizer_weight_decay = self.fake_optimizer_config.get("weight_decay", self.optimizer_weight_decay) - self.fake_optimizer_adam_epsilon = self.fake_optimizer_config.get("adam_epsilon", self.optimizer_adam_epsilon) - - self.dmd_config = parsed.dmd - self.cdm_trick = CdmTrick.from_mapping(self.dmd_config.get("cdm", {})) - self.cdm_trick.validate_trainer( - CdmTrainerConstraints( - supported=self.supports_cdm, - model_name=self.model_config["name"], - supported_model_names=self.supported_cdm_model_names, + dmd_config=self.dmd_config, + student_config=self.student_config, + model_config=self.model_config, + num_inference_steps=self.num_inference_steps, + ) + self.parsed_dmd_schedule_config = parsed + self.num_train_timestep = parsed.num_train_timestep + self.denoising_step_list = parsed.denoising_step_list + self.num_inference_steps = len(self.denoising_step_list) + self.warp_denoising_step = parsed.warp_denoising_step + self.min_step = parsed.min_step + self.max_step = parsed.max_step + self.score_timestep_shift = parsed.score_timestep_shift + self.ts_schedule = parsed.ts_schedule + self.ts_schedule_max = parsed.ts_schedule_max + self.min_score_timestep = parsed.min_score_timestep + self.student_checkpoint_path = parsed.student_checkpoint_path + self.student_checkpoint_strict = parsed.student_checkpoint_strict + + self.diversity_trick = DiversityTrick.from_mapping(self.dmd_config.get("div_loss", {})) + self.diversity_trick.validate_trainer( + DiversityTrainerConstraints( + supported=self.supports_diversity_loss, + rollout_steps=self.num_inference_steps, ), self.trainer_name, ) - self.ida_trick = ImplicitDistributionAlignmentTrick.from_mapping(self.dmd_config.get("ida", {})) - if self.ida_trick.enabled and not self.supports_ida: - raise ValueError(f"{self.trainer_name} does not support training.dmd.ida.") - self.num_inference_steps = parsed.num_inference_steps - self.fake_update_ratio = parsed.fake_update_ratio - self.guidance_scale = parsed.guidance_scale - self.negative_prompt = parsed.negative_prompt - self.cfg_norm = parsed.cfg_norm - self.image_sizes = parsed.image_sizes - self.random_schedule_enabled = parsed.random_schedule_enabled - self.random_schedule_num_steps_min = parsed.random_schedule_num_steps_min - self.random_schedule_num_steps_max = parsed.random_schedule_num_steps_max - self.random_schedule_sigma_min = parsed.random_schedule_sigma_min - self.random_schedule_sigma_max = parsed.random_schedule_sigma_max - self.random_schedule_sampling_method = parsed.random_schedule_sampling_method - - @property - def cdm_enabled(self): - return self.cdm_trick.enabled - - def _get_optimizer_config(self): - return self.training_config["student"]["optimizer"] - - def _setup_trainable_model(self, model, role="student"): - if role == "student": - train_type = self.student_train_type - lora_config = self.student_lora_config - elif role in {"fake", "fake_real"}: - train_type = self.fake_train_type - lora_config = self.fake_lora_config - else: - raise ValueError(f"Unsupported DMD model role: {role}") - if train_type == "lora": - model.add_lora(lora_config["rank"], lora_config["alpha"], lora_config.get("target_modules")) - model.set_lora_trainable() - return - model.set_full_trainable() - - def _restore_trainable_model(self, model, role="student"): - if role == "student": - train_type = self.student_train_type - elif role in {"fake", "fake_real"}: - train_type = self.fake_train_type - else: - raise ValueError(f"Unsupported DMD model role: {role}") - if train_type == "lora": - model.set_lora_trainable() - return - model.set_full_trainable() - - def _save_model_weights(self, model, save_dir, role="student"): - train_type = self.student_train_type if role == "student" else self.fake_train_type - if train_type == "lora": - model.save_lora_weights(save_dir) - return - if is_main_process(): - torch.save(model.denoiser_module().state_dict(), os.path.join(save_dir, "model_state.pt")) - - def _load_model_weights(self, model, save_dir, role="student"): - train_type = self.student_train_type if role == "student" else self.fake_train_type - if train_type == "lora": - model.load_lora_weights_for_resume(save_dir) - return - model_state_path = os.path.join(save_dir, "model_state.pt") - if not os.path.exists(model_state_path): - raise RuntimeError(f"model_state.pt not found in {save_dir}") - state_dict = torch.load(model_state_path, map_location="cpu", weights_only=False) - model.denoiser_module().load_state_dict(state_dict) - - def setup(self, resume_ckpt_path=None): - super().setup(resume_ckpt_path=None) - base_model_config = { - key: copy.deepcopy(value) - for key, value in self.model_config.items() - if key - not in { - "fake", - "teacher", - "student", - "student_2", - "fake_2", - "fake_low_high", - "fake_real", - "fake_real_high", - "fake_real_low", - "teacher_2", - } - } - - fake_model_config = copy.deepcopy(self.config) - fake_model_config["model"] = copy.deepcopy(base_model_config) - if "fake" in self.model_config: - if not isinstance(self.model_config["fake"], dict): - raise ValueError("model.fake must be a mapping.") - fake_model_config["model"].update(copy.deepcopy(self.model_config["fake"])) - self.fake_model_config = copy.deepcopy(fake_model_config) - self.fake_model = build_model(fake_model_config) - self.fake_model.load_components(transformer_only=True, reference_model=self.model) - self._setup_trainable_model(self.fake_model, role="fake") - apply_parallel(self.fake_model, self.config) - if self.gradient_checkpointing: - self.fake_model.enable_gradient_checkpointing() - - teacher_model_config = copy.deepcopy(self.config) - teacher_model_config["model"] = copy.deepcopy(base_model_config) - if "teacher" in self.model_config: - if not isinstance(self.model_config["teacher"], dict): - raise ValueError("model.teacher must be a mapping.") - teacher_model_config["model"].update(copy.deepcopy(self.model_config["teacher"])) - self.teacher_model = build_model(teacher_model_config) - self.teacher_model.load_components(transformer_only=True, reference_model=self.model) - self.teacher_model.transformer.requires_grad_(False) - self.teacher_model.transformer.eval() - apply_parallel(self.teacher_model, self.config) - self.teacher_model.transformer.eval() - - self.fake_trainable_params = list(self.fake_model.trainable_parameters()) - self.fake_optimizer = self._build_optimizer( - self.fake_trainable_params, + diversity_scheduler = ( + DMDFlowMatchingScheduler( + self.config, + self.dmd_config, + ) + if self.diversity_trick.enabled + else None + ) + self.diversity_trick.setup(DiversitySetupContext(scheduler=diversity_scheduler)) + self.real_data_fake_trick = RealDataFakeTrick.from_mappings( { - "learning_rate": self.fake_optimizer_learning_rate, - "adam_beta1": self.fake_optimizer_adam_beta1, - "adam_beta2": self.fake_optimizer_adam_beta2, - "weight_decay": self.fake_optimizer_weight_decay, - "adam_epsilon": self.fake_optimizer_adam_epsilon, - }, - ) - self.fake_lr_scheduler = self._build_lr_scheduler( - self.fake_optimizer, - num_warmup_steps=0, - num_training_steps=max(1, self.max_train_iters * self.fake_update_ratio), + "main": ( + "fake_real", + self.dmd_config.get("fake_real", {}), + ) + } ) + if self.real_data_fake_trick.enabled and not self.supports_real_data_fake: + raise ValueError(f"{self.trainer_name} does not support training.dmd.fake_real.") + self.real_data_fake_trick.validate_schedule( + mode="standard", + num_train_timestep=self.num_train_timestep, + ) + if self.real_data_fake_trick.enabled: + train_data_name = self.config["data"]["train"].get("name") + if train_data_name not in { + "latent_dataset", + "video_dataset", + }: + raise ValueError("Real-data fake training requires data.train.name=latent_dataset or video_dataset.") + self.fake_real_train_type = self.fake_train_type + self.fake_real_lora_config = copy.deepcopy(self.fake_lora_config) + self.fake_real_optimizer_config = copy.deepcopy(self.fake_optimizer_config) - self.scheduler = DMDFlowMatchingScheduler(self.config, self.dmd_config) - + def setup(self, resume_ckpt_path=None): + if resume_ckpt_path is None and self.student_checkpoint_path: + self._load_student_checkpoint(self.student_checkpoint_path, strict=self.student_checkpoint_strict) + super().setup(resume_ckpt_path=None) + self._setup_fake_real_resources() if resume_ckpt_path is not None: self._load_resume_state(resume_ckpt_path) - - if not self.defer_ida_setup: - self._setup_ida_trick() - - logger.info("[train] dmd student model={} path={}", self.model_config["name"], self.model_config["pretrained_model_name_or_path"]) - logger.info("[train] dmd fake model={} path={}", fake_model_config["model"]["name"], fake_model_config["model"]["pretrained_model_name_or_path"]) - logger.info("[train] dmd teacher model={} path={}", teacher_model_config["model"]["name"], teacher_model_config["model"]["pretrained_model_name_or_path"]) - logger.info("[train] dmd train_types student={} fake={}", self.student_train_type, self.fake_train_type) - logger.info("[train] dmd student trainable params={}", self._count_trainable(self.model.transformer)) - logger.info("[train] dmd fake trainable params={}", self._count_trainable(self.fake_model.transformer)) - if self.random_schedule_enabled: - logger.info( - "[train] dmd random sigma schedule enabled: steps=[{}, {}], sigma=[{}, {}], sampling_method={}", - self.random_schedule_num_steps_min, - self.random_schedule_num_steps_max, - self.random_schedule_sigma_min, - self.random_schedule_sigma_max, - self.random_schedule_sampling_method, + if resume_ckpt_path is None: + self.lr_scheduler = self._build_lr_scheduler( + self.optimizer, + num_training_steps=max(1, self.max_train_iters), ) - if self.cdm_enabled: - logger.info( - "[train] dmd CDM enabled: weight={} warmup_iters={}", - self.cdm_trick.config.weight, - self.cdm_trick.config.warmup_iters, + self.fake_lr_scheduler = self._build_lr_scheduler( + self.fake_optimizer, + num_warmup_steps=0, + num_training_steps=max( + 1, + self.max_train_iters * self.fake_update_ratio, + ), ) - @staticmethod - def _count_trainable(module): - return sum(1 for param in module.parameters() if param.requires_grad) - - def _ida_model_pairs(self): - return { - "main": IdaModelPair( - student=self.model.denoiser_module(), - fake=self.fake_model.denoiser_module(), - ) - } - - def _setup_ida_trick(self): - self.ida_trick.setup(IdaSetupContext(model_pairs=self._ida_model_pairs())) + time_shift_settings = self.config["scheduler"].get("time_shift_settings", {}) + self.denoising_scheduler = CausalForcingFlowMatchScheduler( + num_train_timesteps=self.config["scheduler"].get("num_train_timesteps", 1000), + time_shift_settings=time_shift_settings, + ) + self.denoising_steps = self._build_denoising_steps(self.student.device) + self.denoising_sigmas = (self.denoising_steps / self.num_train_timestep).to(dtype=torch.float32) + self.real_data_fake_trick.setup(self._real_data_fake_setup_context()) logger.info( - "[train] {} SenseFlow IDA enabled={} decay={}", + "[train] {} denoising_steps={} warped={}", self.trainer_name, - self.ida_trick.enabled, - self.ida_trick.config.decay, + [round(float(step), 4) for step in self.denoising_steps.detach().cpu()], + self.warp_denoising_step, ) - def _after_student_optimizer_step(self, role): - self.ida_trick.after_student_step(IdaStepContext(role=role)) - - @staticmethod - def _do_cfg(cond_pred, uncond_pred, cfg_scale, cfg_norm): - return do_cfg( - cond_pred, - uncond_pred, - cfg_scale, - cfg_norm, + def _setup_fake_real_resources(self): + self.fake_real_model = None + self.fake_real_optimizer = None + self.fake_real_lr_scheduler = None + self.fake_real_trainable_params = [] + if not self.real_data_fake_trick.enabled_for("main"): + return + model_config = copy.deepcopy(self.fake_model_config) + self.fake_real_model = build_loaded_model( + model_config, + transformer_only=True, + reference_model=self.model, ) - - @staticmethod - def _dmd_loss(latents, x_pred_fake_flow, x_pred_teacher, norm_clip_min=None): - return dmd_loss( - latents, - x_pred_fake_flow, - x_pred_teacher, - norm_clip_min=norm_clip_min, + self.fake_real = self.fake_real_model.capabilities.require(DistillationCapability) + self._setup_trainable_model( + self.fake_real_model, + role="fake_real", ) - - def _prepare_sampling_schedule(self, latent_shape): - latent_hw = latent_shape[-2:] - if self.random_schedule_enabled: - num_steps = self._sample_synced_int(self.random_schedule_num_steps_min, self.random_schedule_num_steps_max + 1) - self.scheduler.set_random_timesteps( - self.random_schedule_num_steps_min, - self.random_schedule_num_steps_max, - sigma_min=self.random_schedule_sigma_min, - sigma_max=self.random_schedule_sigma_max, - sampling_method=self.random_schedule_sampling_method, - latent_hw=latent_hw, - device=self.model.device, - num_steps=num_steps, - ) - return - self.scheduler.set_timesteps(self.num_inference_steps, latent_hw=latent_hw, device=self.model.device) - - def _latent_shape(self, sample): - image = sample["inputs"].get("target_image") - if image is None: - raise KeyError("DMD image latent shape expects inputs.target_image.") - batch_size = image.shape[0] - if self.image_sizes: - image_size_index = int(torch.randint(0, len(self.image_sizes), (1,), device=self.model.device).item()) - image_size_index = broadcast_sequence_parallel_value(image_size_index) - height, width = self.image_sizes[image_size_index] - else: - height, width = image.shape[-2], image.shape[-1] - - latent_channels = getattr(self.model.vae.config, "z_dim", None) - if latent_channels is None: - latent_channels = self.model.transformer.config.in_channels // 4 - return ( - batch_size, - int(latent_channels), - 1, - height // self.model.vae_scale_factor, - width // self.model.vae_scale_factor, - ) - - def _encode_conditions(self, sample): - conditioning = sample["conditioning"] - prompt = conditioning.get("prompt", "") - with torch.no_grad(): - condition = self.model.encode_prompt_condition(prompt) - if self.guidance_scale > 1: - is_scalar_prompt = isinstance(prompt, str) - batch_size = 1 if is_scalar_prompt else len(prompt) - negative_prompt = self._negative_prompt_for_conditioning( - conditioning, - batch_size, - return_scalar=is_scalar_prompt, - ) - negative_condition = self.model.encode_prompt_condition(negative_prompt) - else: - negative_condition = None - condition = broadcast_sequence_parallel_value(condition) - negative_condition = broadcast_sequence_parallel_value(negative_condition) if negative_condition is not None else None - return condition, negative_condition - - def _negative_prompt_for_conditioning(self, conditioning, batch_size, return_scalar=False): - negative_prompt = conditioning.get("negative_prompt") - if negative_prompt is None: - values = [] - elif isinstance(negative_prompt, str): - values = [negative_prompt] - else: - values = list(negative_prompt) - - if not values: - values = [self.negative_prompt] * batch_size - elif len(values) == 1 and batch_size > 1: - values *= batch_size - elif len(values) != batch_size: - raise ValueError(f"Expected {batch_size} negative prompts, got {len(values)}.") - - values = [value if isinstance(value, str) and value.strip() else self.negative_prompt for value in values] - return values[0] if return_scalar else values - - def _predict_velocity(self, model, latents, sigma, condition): - denoiser_input = model.prepare_denoiser_input(latents) - prediction = model.denoise(denoiser_input, sigma, condition) - prediction = model.postprocess_denoiser_output(prediction, denoiser_input) - return prediction - - def _predict_teacher_velocity(self, latents, sigma, condition, negative_condition): - if negative_condition is None: - return self._predict_velocity(self.teacher_model, latents, sigma, condition) - - if self.teacher_model.cfg_on_denoiser_output(): - denoiser_input = self.teacher_model.prepare_denoiser_input(latents) - cond_prediction = self.teacher_model.denoise(denoiser_input, sigma, condition) - uncond_prediction = self.teacher_model.denoise(denoiser_input, sigma, negative_condition) - prediction = self._do_cfg(cond_prediction, uncond_prediction, self.guidance_scale, self.cfg_norm) - return self.teacher_model.postprocess_denoiser_output(prediction, denoiser_input) - - velocity_teacher_cond = self._predict_velocity(self.teacher_model, latents, sigma, condition) - velocity_teacher_uncond = self._predict_velocity(self.teacher_model, latents, sigma, negative_condition) - return self._do_cfg(velocity_teacher_cond, velocity_teacher_uncond, self.guidance_scale, self.cfg_norm) - - def sample_initial_latents(self, latent_shape): - return broadcast_sequence_parallel_value(torch.randn(latent_shape, device=self.model.device, dtype=self.running_dtype)) - - def _sample_synced_int(self, low, high): - value = torch.randint(int(low), int(high), (1,), device=self.model.device, dtype=torch.int64) - if is_distributed(): - dist.broadcast(value, src=0) - return int(value.item()) - - def sample_end_step(self): - return self._sample_synced_int(0, self.scheduler.num_inference_steps) - - def run_back_simulation(self, condition, latent_shape, end_step_idx, grad_enabled, xt=None): - if xt is None: - xt = self.sample_initial_latents(latent_shape) - x0 = None - xt_end = None - vt_end = None - self.model.transformer.train() - for idx in range(end_step_idx + 1): - sigma = self.scheduler.sigma_at(idx, latent_shape[0], device=self.model.device, dtype=self.running_dtype) - context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad - with context(): - velocity = self._predict_velocity(self.model, xt, sigma, condition) - if idx == end_step_idx: - xt_end = xt.detach() - vt_end = velocity.detach() - xt, x0 = self.scheduler.step_by_index(velocity, idx, xt) - return x0, xt_end, vt_end - - def _cdm_step_context( - self, - xt, - vt, - end_step_idx, - condition, - current_iter, - ): - return CdmStepContext( - trajectory_latent=xt, - trajectory_velocity=vt, - end_step_index=end_step_idx, - condition=condition, - current_iteration=current_iter, - device=self.model.device, - dtype=self.running_dtype, - scheduler=self.scheduler, - predict_student_velocity=lambda latent, sigma, cond: ( - self._predict_velocity( - self.model, - latent, - sigma, - cond, - ) - ), - predict_fake_velocity=lambda latent, sigma, cond: ( - self._predict_velocity( - self.fake_model, - latent, - sigma, - cond, - ) - ), - predict_teacher_velocity=lambda latent, sigma, cond: ( - self._predict_teacher_velocity( - latent, - sigma, - cond, - None, - ) + self.fake_real_model.capabilities.require(ParallelCapability).apply(self.config) + if self.gradient_checkpointing: + self.fake_real_model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() + self.fake_real_trainable_params = list(self.fake_real_model.capabilities.require(TrainableModelCapability).parameters()) + self.fake_real_optimizer = self._build_optimizer( + self.fake_real_trainable_params, + self.fake_real_optimizer_config, + ) + self.fake_real_lr_scheduler = self._build_lr_scheduler( + self.fake_real_optimizer, + num_warmup_steps=0, + num_training_steps=max( + 1, + self.max_train_iters * self.fake_update_ratio, ), - prepare_fake_model=self.fake_model.transformer.eval, - dmd_loss=self._dmd_loss, + ) + logger.info( + "[train] {} independent fake_real model={} path={} train_type={} trainable_params={}", + self.trainer_name, + model_config["model"]["name"], + model_config["model"]["pretrained_model_name_or_path"], + self.fake_real_train_type, + len(self.fake_real_trainable_params), ) - def forward_loss(self, latent_shape, conditions, stage, current_iter=None): - condition, negative_condition = conditions - self._prepare_sampling_schedule(latent_shape) - end_step_idx = self.sample_end_step() - xt_start = self.sample_initial_latents(latent_shape) - x0, xt_end, vt_end = self.run_back_simulation(condition, latent_shape, end_step_idx, grad_enabled=(stage != "fake"), xt=xt_start) - - sigma = self.scheduler.sample_renoise_sigma(latent_shape[0], device=self.model.device, dtype=self.running_dtype) - sigma = broadcast_sequence_parallel_value(sigma) - noise = broadcast_sequence_parallel_value(torch.randn(latent_shape, device=self.model.device, dtype=torch.float32)) - renoised_xt = self.scheduler.add_noise(x0.detach(), noise, sigma) - - if stage == "fake": - self.fake_model.transformer.train() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition) - velocity_gt = self.scheduler.build_train_gt(x0.float(), noise) - loss_fake = F.mse_loss(velocity_fake.float(), velocity_gt.float(), reduction="mean") - return {"fake": loss_fake} - - with torch.no_grad(): - self.fake_model.transformer.eval() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition) - velocity_teacher = self._predict_teacher_velocity(renoised_xt, sigma, condition, negative_condition) - - expanded_sigma = self.scheduler._expand_to_ndim(sigma, renoised_xt.ndim) - x_pred_fake = renoised_xt - expanded_sigma * velocity_fake - x_pred_teacher = renoised_xt - expanded_sigma * velocity_teacher - loss_dmd = self._dmd_loss(x0, x_pred_fake, x_pred_teacher) - cdm_result = self.cdm_trick.student_loss( - self._cdm_step_context( - xt_end, - vt_end, - end_step_idx, - condition, - current_iter, - ) + def _load_student_checkpoint(self, checkpoint_path, strict=True): + model_state_path = checkpoint_path + if os.path.isdir(model_state_path): + dist_state_path = os.path.join(model_state_path, "dist_state") + if os.path.isdir(dist_state_path): + module = self.parallel.state_module() + options = StateDictOptions(ignore_frozen_params=False, strict=strict) + model_state = get_model_state_dict(module, options=options) + state = {"model": model_state} + dcp.load(state, checkpoint_id=dist_state_path) + set_model_state_dict( + module, + model_state_dict=state["model"], + options=options, + ) + logger.info("[train] loaded {} student checkpoint from distributed state {}", self.trainer_name, dist_state_path) + return + model_state_path = os.path.join(model_state_path, "model_state.pt") + if not os.path.exists(model_state_path): + raise RuntimeError(f"{self.trainer_name} student checkpoint not found: {checkpoint_path}") + + state = torch.load(model_state_path, map_location="cpu", weights_only=False) + for key in ("generator_ema", "generator", "model", "state_dict"): + if isinstance(state, dict) and key in state: + state = state[key] + break + + fixed = {} + for key, value in state.items(): + for prefix in ( + "model._fsdp_wrapped_module.", + "model._checkpoint_wrapped_module.", + "model._orig_mod.", + "model.", + "_fsdp_wrapped_module.", + "_checkpoint_wrapped_module.", + "_orig_mod.", + ): + if key.startswith(prefix): + key = key[len(prefix) :] + fixed[key] = value + + incompatible = self.parallel.state_module().load_state_dict( + fixed, + strict=strict, ) - return { - "student": loss_dmd + cdm_result.loss, - "dmd": loss_dmd.detach(), - "cdm": cdm_result.metrics["cdm"], - "cdm_weight": cdm_result.metrics["cdm_weight"], - } + if not strict: + if incompatible.missing_keys: + logger.warning("Missing keys when loading {} student checkpoint: {}", self.trainer_name, incompatible.missing_keys) + if incompatible.unexpected_keys: + logger.warning("Unexpected keys when loading {} student checkpoint: {}", self.trainer_name, incompatible.unexpected_keys) + logger.info("[train] loaded {} student checkpoint path={}", self.trainer_name, model_state_path) def train(self): resume_ckpt_path, current_iter = self._resolve_resume() @@ -530,14 +255,12 @@ def train(self): max_train_iters = self.max_train_iters grad_accum_iters = max(1, int(self.gradient_accumulation_iters)) - fake_update_ratio = self.fake_update_ratio - max_grad_norm = self.max_grad_norm save_every_iters = self.save_every_iters save_total_limit = self.save_total_limit logger.info( "[train] start method={} student_train_type={} fake_train_type={} iter={}/{} world_size={} grad_accum={} train_log_every_iters={} fake_update_ratio={}", - self.training_config.get("method", "dmd"), + self.training_config.get("method", self.trainer_name), self.student_train_type, self.fake_train_type, current_iter, @@ -545,7 +268,23 @@ def train(self): get_world_size(), grad_accum_iters, self.train_log_every_iters, - fake_update_ratio, + self.fake_update_ratio, + ) + logger.info( + "[train] {} diversity enabled={} weight={} teacher_steps={} anchor_step={}", + self.trainer_name, + self.diversity_trick.enabled, + self.diversity_trick.config.weight, + self.diversity_trick.config.teacher_inference_steps, + self.diversity_trick.config.anchor_step, + ) + real_data_config = self.real_data_fake_trick.config.regions["main"] + logger.info( + "[train] {} real_data_fake enabled={} weight={} timesteps={}", + self.trainer_name, + real_data_config.enabled, + real_data_config.weight, + list(real_data_config.timestep_list), ) if self.infer_every_iters: self.inferencer.set_data(self.dataloader_eval) @@ -554,101 +293,65 @@ def train(self): samples = self._iter_train_samples() while current_iter < max_train_iters: - self.optimizer.zero_grad(set_to_none=True) - running_dmd = 0.0 - running_cdm = 0.0 - running_cdm_weight = 0.0 - for micro_idx in range(grad_accum_iters): - sample = next(samples) - conditions = self._encode_conditions(sample) - latent_shape = self._latent_shape(sample) - self._set_student_gradient_sync(micro_idx == grad_accum_iters - 1) - res_student = self.forward_loss(latent_shape, conditions, stage="student", current_iter=current_iter) - loss_student = res_student["student"] - (loss_student / grad_accum_iters).backward() - running_dmd += res_student["dmd"].item() / grad_accum_iters - if self.cdm_enabled: - running_cdm += res_student["cdm"].item() / grad_accum_iters - running_cdm_weight = res_student["cdm_weight"] - - self._sync_sequence_parallel_grads(self.trainable_params) - torch.nn.utils.clip_grad_norm_(self.trainable_params, max_grad_norm) - self.optimizer.step() - self._after_student_optimizer_step("main") - self.lr_scheduler.step() - self.optimizer.zero_grad(set_to_none=True) - - running_fake = 0.0 - for _ in range(fake_update_ratio): - self.fake_optimizer.zero_grad(set_to_none=True) - fake_loss = 0.0 - for micro_idx in range(grad_accum_iters): - sample = next(samples) - conditions = self._encode_conditions(sample) - latent_shape = self._latent_shape(sample) - self._set_fake_gradient_sync(micro_idx == grad_accum_iters - 1) - res_fake = self.forward_loss( - latent_shape, - conditions, - stage="fake", - ) - loss_fake = res_fake["fake"] - (loss_fake / grad_accum_iters).backward() - fake_loss += loss_fake.item() / grad_accum_iters - self._sync_sequence_parallel_grads(self.fake_trainable_params) - torch.nn.utils.clip_grad_norm_( - self.fake_trainable_params, - max_grad_norm, + student_result = self._train_one_stage( + samples, + stage="student", + grad_accum_iters=grad_accum_iters, + current_iter=current_iter, + ) + loss_dmd_value = student_result["dmd"] + loss_cdm_value = student_result.get("cdm", 0.0) + cdm_weight = student_result.get("cdm_weight", 0.0) + loss_div_value = student_result["div_loss"] + loss_real_dmd_value = student_result["real_dmd"] + loss_fake_value = 0.0 + loss_fake_real_value = 0.0 + for _ in range(self.fake_update_ratio): + fake_result = self._train_one_stage( + samples, + stage="fake", + grad_accum_iters=grad_accum_iters, + current_iter=current_iter, ) - self.fake_optimizer.step() - self.fake_lr_scheduler.step() - self.fake_optimizer.zero_grad(set_to_none=True) - running_fake += fake_loss / fake_update_ratio + loss_fake_value += fake_result["loss"] + loss_fake_real_value += fake_result["fake_real"] + loss_fake_value /= self.fake_update_ratio + loss_fake_real_value /= self.fake_update_ratio current_iter += 1 - display_dmd = reduce_mean(running_dmd) - display_fake = reduce_mean(running_fake) + display_fake = reduce_mean(loss_fake_value) + display_dmd = reduce_mean(loss_dmd_value) if loss_dmd_value is not None else None + display_cdm = reduce_mean(loss_cdm_value) + display_div = reduce_mean(loss_div_value) + display_real_dmd = reduce_mean(loss_real_dmd_value) + display_fake_real = reduce_mean(loss_fake_real_value) current_lr = self.lr_scheduler.get_last_lr()[0] if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: + dmd_text = "nan" if display_dmd is None else f"{display_dmd:.6f}" + logger.info( + "[train] iter={}/{} dmd={} div_loss={:.6f} real_dmd={:.6f} fake={:.6f} fake_real={:.6f} lr={:.8f}", + current_iter, + max_train_iters, + dmd_text, + display_div, + display_real_dmd, + display_fake, + display_fake_real, + current_lr, + ) + metrics = { + "train/div_loss": display_div, + "train/real_dmd": display_real_dmd, + "train/fake": display_fake, + "train/fake_real": display_fake_real, + "train/lr": current_lr, + } + if display_dmd is not None: + metrics["train/dmd"] = display_dmd if self.cdm_enabled: - display_cdm = reduce_mean(running_cdm) - logger.info( - "[train] iter={}/{} dmd={:.6f} cdm={:.6f} cdm_w={:.6f} fake={:.6f} lr={:.8f}", - current_iter, - max_train_iters, - display_dmd, - display_cdm, - running_cdm_weight, - display_fake, - current_lr, - ) - self.log_metrics( - { - "train/dmd": display_dmd, - "train/cdm": display_cdm, - "train/cdm_weight": running_cdm_weight, - "train/fake": display_fake, - "train/lr": current_lr, - }, - step=current_iter, - ) - else: - logger.info( - "[train] iter={}/{} dmd={:.6f} fake={:.6f} lr={:.8f}", - current_iter, - max_train_iters, - display_dmd, - display_fake, - current_lr, - ) - self.log_metrics( - { - "train/dmd": display_dmd, - "train/fake": display_fake, - "train/lr": current_lr, - }, - step=current_iter, - ) + metrics["train/cdm"] = display_cdm + metrics["train/cdm_weight"] = cdm_weight + self.log_metrics(metrics, step=current_iter) if save_every_iters and current_iter % save_every_iters == 0: self.save_checkpoint(current_iter, save_total_limit) @@ -658,58 +361,462 @@ def train(self): logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) - def _iter_train_samples(self): - epoch = 0 - while True: - sampler = getattr(self.dataloader_train, "sampler", None) - if hasattr(sampler, "set_epoch"): - sampler.set_epoch(epoch) - for sample in self.dataloader_train: - yield sample - epoch += 1 + def _train_one_stage( + self, + samples, + stage, + grad_accum_iters, + current_iter=None, + ): + if stage == "student": + optimizer = self.optimizer + scheduler = self.lr_scheduler + params = self.trainable_params + set_sync = self._set_student_gradient_sync + elif stage == "fake": + optimizer = self.fake_optimizer + scheduler = self.fake_lr_scheduler + params = self.fake_trainable_params + set_sync = self._set_fake_gradient_sync + else: + raise ValueError(f"Unsupported {self.trainer_name} training stage: {stage}") + + optimizer.zero_grad(set_to_none=True) + loss_value = 0.0 + use_real_data_fake = self.real_data_fake_trick.enabled_for("main") + if stage == "fake" and use_real_data_fake: + self.fake_real_optimizer.zero_grad(set_to_none=True) + if stage == "student": + metric_values = { + "div_loss": 0.0, + "real_dmd": 0.0, + } + else: + metric_values = {"fake_real": 0.0} + for micro_idx in range(grad_accum_iters): + sample = next(samples) + conditions = self._encode_conditions(sample) + latent_shape = self._latent_shape(sample) + initial_noise = self.sample_initial_latents(latent_shape) + set_sync((stage == "student" and use_real_data_fake) or micro_idx == grad_accum_iters - 1) + result = self.forward_loss( + latent_shape, + conditions, + stage=stage, + initial_noise=initial_noise, + current_iter=current_iter, + ) + if isinstance(result, dict): + loss = result["loss"] + for name, value in result.items(): + if name == "loss": + continue + scalar = value.item() if torch.is_tensor(value) else float(value) + metric_values[name] = metric_values.get(name, 0.0) + scalar / grad_accum_iters + del result + else: + loss = result + (loss / grad_accum_iters).backward() + loss_value += loss.item() / grad_accum_iters + del loss + if stage == "student" and self.diversity_trick.enabled: + div_raw, _ = self._backward_diversity_loss( + initial_noise, + conditions, + grad_accum_iters, + ) + metric_values["div_loss"] += div_raw + del initial_noise + if stage == "student" and use_real_data_fake: + real_result = self.real_data_fake_trick.student_loss( + self._real_data_fake_context( + sample, + conditions, + ) + ) + (real_result.loss / grad_accum_iters).backward() + metric_values["real_dmd"] += real_result.metrics["real_dmd"].item() / grad_accum_iters + del real_result + elif stage == "fake" and use_real_data_fake: + self._set_fake_real_gradient_sync(micro_idx == grad_accum_iters - 1) + fake_real_result = self.real_data_fake_trick.fake_loss( + self._real_data_fake_context( + sample, + conditions, + ) + ) + (fake_real_result.loss / grad_accum_iters).backward() + metric_values["fake_real"] += fake_real_result.metrics["fake_real"].item() / grad_accum_iters + del fake_real_result + + self._sync_sequence_parallel_grads(params) + torch.nn.utils.clip_grad_norm_(params, self.max_grad_norm) + optimizer.step() + if stage == "student": + self._after_student_optimizer_step("main") + scheduler.step() + optimizer.zero_grad(set_to_none=True) + if stage == "fake" and use_real_data_fake: + self._sync_sequence_parallel_grads(self.fake_real_trainable_params) + torch.nn.utils.clip_grad_norm_( + self.fake_real_trainable_params, + self.max_grad_norm, + ) + self.fake_real_optimizer.step() + self.fake_real_lr_scheduler.step() + self.fake_real_optimizer.zero_grad(set_to_none=True) + if metric_values: + return { + "loss": loss_value, + **metric_values, + } + return loss_value + + def _build_denoising_steps(self, device): + raw_steps = torch.tensor(self.denoising_step_list, dtype=torch.long, device=device) + if not self.warp_denoising_step: + return raw_steps.to(dtype=torch.float32) + + timesteps = torch.cat( + [ + self.denoising_scheduler.timesteps.to(device=device, dtype=torch.float32), + torch.zeros(1, device=device, dtype=torch.float32), + ] + ) + indices = self.denoising_scheduler.num_train_timesteps - raw_steps + return timesteps[indices] + + def _extract_real_latents(self, sample): + return self.student.extract_real_latents( + sample, + self.running_dtype, + broadcast_sequence_parallel_value, + ) + + def _predict_real_student_velocity( + self, + region, + xt, + sigma, + condition, + grad_enabled, + ): + del region + context = torch.enable_grad if grad_enabled else torch.no_grad + self.student.set_training(grad_enabled) + with context(): + return self._predict_velocity( + self.student, + xt, + sigma, + condition, + ) - def _set_student_gradient_sync(self, enabled): - set_parallel_gradient_sync(self.model, enabled) + def _predict_real_teacher_velocity( + self, + region, + xt, + sigma, + condition, + negative_condition, + ): + del region + self.teacher.set_training(False) + return self._predict_teacher_velocity( + xt, + sigma, + condition, + negative_condition, + ) - def _set_fake_gradient_sync(self, enabled): - set_parallel_gradient_sync(self.fake_model, enabled) + def _diversity_step_context(self, initial_noise, conditions): + return DiversityStepContext( + initial_noise=initial_noise.detach(), + condition=conditions[0], + negative_condition=conditions[1], + latent_hw=initial_noise.shape[-2:], + device=self.student.device, + dtype=self.running_dtype, + predict_teacher_velocity=self._predict_teacher_velocity, + predict_student_velocity=partial( + self._predict_velocity, + self.student, + ), + student_scheduler=self.scheduler, + expand_to_ndim=self.scheduler._expand_to_ndim, + ) + + def _backward_diversity_loss( + self, + initial_noise, + conditions, + grad_accum_iters, + ): + self.diversity_trick.prepare_models( + self.teacher.denoiser(), + self.student.denoiser(), + ) + result = self.diversity_trick.student_loss(self._diversity_step_context(initial_noise, conditions)) + (result.loss / grad_accum_iters).backward() + raw_value = result.metrics["div_loss"].item() / grad_accum_iters + weighted_value = result.loss.detach().item() / grad_accum_iters + return raw_value, weighted_value + + def _real_data_fake_setup_context(self): + return RealDataFakeSetupContext( + mode="standard", + scheduler=self.scheduler, + denoising_scheduler=self.denoising_scheduler, + num_train_timestep=self.num_train_timestep, + warp_denoising_step=self.warp_denoising_step, + score_timestep_shift=self.score_timestep_shift, + min_step=self.min_step, + max_step=self.max_step, + min_score_timestep=self.min_score_timestep, + ) + + def _real_data_fake_context(self, sample, conditions, region="main"): + return RealDataFakeStepContext( + region=region, + fake_model=self.fake_real, + sample=sample, + condition=conditions[0], + negative_condition=conditions[1], + device=self.student.device, + dtype=self.running_dtype, + extract_real_latents=self._extract_real_latents, + sample_synced_int=self._sample_synced_int, + broadcast_noise=broadcast_sequence_parallel_value, + predict_student_velocity=self._predict_real_student_velocity, + predict_fake_velocity=self._predict_velocity, + predict_teacher_velocity=(self._predict_real_teacher_velocity), + dmd_loss=self._dmd_loss, + ) + + def _set_fake_real_gradient_sync(self, enabled): + if self.fake_real_model is not None: + self.fake_real_model.capabilities.require(ParallelCapability).set_gradient_sync(enabled) def _set_gradient_sync(self, enabled): - self._set_student_gradient_sync(enabled) - self._set_fake_gradient_sync(enabled) + super()._set_gradient_sync(enabled) + self._set_fake_real_gradient_sync(enabled) - def _sync_sequence_parallel_grads(self, params): - sync_sequence_parallel_gradients(params) + def forward_loss( + self, + latent_shape, + conditions, + stage, + initial_noise=None, + current_iter=None, + ): + condition, negative_condition = conditions + ( + generated, + denoised_timestep_from, + denoised_timestep_to, + ) = self.run_back_simulation( + condition, + latent_shape, + grad_enabled=(stage != "fake"), + xt=initial_noise, + ) - def _fake_weights_dir(self, root_dir): - return self.checkpoint_manager._fake_weights_dir(root_dir) + sigma = self._sample_score_sigma( + denoised_timestep_from=denoised_timestep_from, + denoised_timestep_to=denoised_timestep_to, + device=self.student.device, + dtype=self.running_dtype, + ) + noise = self.student.random_noise_like( + generated, + torch.float32, + broadcast_sequence_parallel_value, + ) - def _trick_checkpoint_metadata(self): - return self.checkpoint_manager._trick_checkpoint_metadata() + with torch.no_grad(): + renoised_xt = self.student.add_noise( + self.scheduler, + generated, + noise, + sigma, + ) - def _validate_optional_trick_metadata(self, state, state_path): - return self.checkpoint_manager._validate_optional_trick_metadata(state, state_path) + if stage == "fake": + self.fake.set_training(False) + velocity_fake = self._predict_velocity( + self.fake, + renoised_xt, + sigma, + condition, + ) + velocity_gt = self.student.training_target(generated, noise) + return self.student.regression_loss( + velocity_fake, + velocity_gt, + ) - def _load_resume_state(self, resume_ckpt_path): - return self.checkpoint_manager._load_resume_state(resume_ckpt_path) + with torch.no_grad(): + self.fake.set_training(False) + self.teacher.set_training(False) + velocity_fake = self._predict_velocity( + self.fake, + renoised_xt, + sigma, + condition, + ) + velocity_teacher = self.teacher.predict_guided_velocity( + renoised_xt, + sigma, + condition, + negative_condition if self.guidance_scale != 0 else None, + self.guidance_scale, + self.cfg_norm, + ) - def _validate_dmd_checkpoint_metadata(self, state, state_path, resume_ckpt_path): - return self.checkpoint_manager._validate_dmd_checkpoint_metadata(state, state_path, resume_ckpt_path) + x_pred_fake = self.student.x0_from_velocity( + renoised_xt, + velocity_fake, + sigma, + ) + x_pred_teacher = self.student.x0_from_velocity( + renoised_xt, + velocity_teacher, + sigma, + ) - def _load_single_process_state(self, resume_ckpt_path): - return self.checkpoint_manager._load_single_process_state(resume_ckpt_path) + loss_dmd = self.student.dmd_loss( + generated, + x_pred_fake, + x_pred_teacher, + ) + result = { + "loss": loss_dmd, + "dmd": loss_dmd.detach(), + } + if not self.cdm_enabled: + return result + + end_sample, end_velocity, end_step_index = self._last_rollout_end + cdm_result = self.cdm_trick.student_loss( + self._cdm_step_context( + end_sample, + end_velocity, + end_step_index, + condition, + current_iter, + ) + ) + result.update( + loss=loss_dmd + cdm_result.loss, + cdm=cdm_result.metrics["cdm"], + cdm_weight=cdm_result.metrics["cdm_weight"], + ) + return result - def _load_distributed_state(self, resume_ckpt_path): - return self.checkpoint_manager._load_distributed_state(resume_ckpt_path) + def sample_end_step(self): + return self._sample_synced_int( + self.diversity_trick.minimum_dmd_step_index, + self.scheduler.num_inference_steps, + ) - def save_checkpoint(self, iteration, save_total_limit): - return self.checkpoint_manager.save_checkpoint(iteration, save_total_limit) + def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): + if self.random_schedule_enabled: + self._prepare_sampling_schedule(latent_shape) + self._active_denoising_steps = self.scheduler.timesteps + else: + self.scheduler.set_timesteps( + self.num_inference_steps, + sigmas=[float(sigma) for sigma in self.denoising_sigmas.detach().cpu()], + latent_hw=self.student.latent_hw(latent_shape), + device=self.student.device, + ) + self._active_denoising_steps = self.denoising_steps + if xt is None: + xt = self.sample_initial_latents(latent_shape) - def _should_save_consolidated_student(self): - return self.checkpoint_manager._should_save_consolidated_student() + end_step_idx = ( + self.sample_end_step() + if grad_enabled + else self._sample_synced_int( + 0, + self.scheduler.num_inference_steps, + ) + ) + x0 = None + end_sample = None + end_velocity = None + # Zoe runs every DMD role in eval mode. Gradients are controlled by + # autograd contexts, not by module.train(), which keeps rollout + # dropout and other training-only behavior deterministic. + self.student.set_training(False) + for idx in range(end_step_idx + 1): + sigma = self.scheduler.sigma_at( + idx, + device=self.student.device, + dtype=self.running_dtype, + ) + context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad + with context(): + velocity = self._predict_velocity( + self.student, + xt, + sigma, + condition, + ) + if grad_enabled and idx == end_step_idx: + end_sample = self.student.detach(xt) + end_velocity = self.student.detach(velocity) + xt, x0 = self.student.step( + self.scheduler, + velocity, + idx, + xt, + ) - def _save_consolidated_student_weights(self, save_dir): - return self.checkpoint_manager._save_consolidated_student_weights(save_dir) + if grad_enabled: + self._last_rollout_end = ( + end_sample, + end_velocity, + end_step_idx, + ) - def _save_distributed_state(self, save_dir, iteration): - return self.checkpoint_manager._save_distributed_state(save_dir, iteration) + return self.student.to_dtype( + x0, + self.running_dtype, + ), *self._denoised_timestep_window(end_step_idx) + + def _sample_score_sigma(self, denoised_timestep_from, denoised_timestep_to, device, dtype): + min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep + max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep + min_timestep = max(0, int(min_timestep)) + max_timestep = min(self.num_train_timestep, int(max_timestep)) + if max_timestep <= min_timestep: + max_timestep = min(self.num_train_timestep, min_timestep + 1) + + timestep = torch.randint(min_timestep, max_timestep, (1,), device=device, dtype=torch.long).float() + if self.score_timestep_shift > 1: + t = timestep / self.num_train_timestep + timestep = self.score_timestep_shift * t / (1 + (self.score_timestep_shift - 1) * t) * self.num_train_timestep + timestep = timestep.clamp(self.min_step, self.max_step) + return broadcast_sequence_parallel_value((timestep / self.num_train_timestep).to(dtype=dtype)) + + def _denoised_timestep_window(self, exit_idx): + exit_idx = int(exit_idx) + steps = getattr( + self, + "_active_denoising_steps", + self.denoising_steps, + ) + denoised_timestep_from = self._raw_timestep_from_warped_step(steps[exit_idx]) + if exit_idx == len(steps) - 1: + denoised_timestep_to = 0 + else: + denoised_timestep_to = self._raw_timestep_from_warped_step(steps[exit_idx + 1]) + return denoised_timestep_from, denoised_timestep_to + + def _raw_timestep_from_warped_step(self, warped_step): + if not self.warp_denoising_step: + return int(round(float(warped_step))) + timesteps = self.denoising_scheduler.timesteps.to(device=warped_step.device, dtype=torch.float32) + index = torch.argmin((timesteps - warped_step.float()).abs(), dim=0).item() + return self.denoising_scheduler.num_train_timesteps - int(index) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/video_ar_trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/video_ar_trainer.py deleted file mode 100644 index feadb459f..000000000 --- a/lightx2v_train/lightx2v_train/trainers/dmd/video_ar_trainer.py +++ /dev/null @@ -1,253 +0,0 @@ -import torch -import torch.distributed as dist - -from lightx2v_train.runtime.distributed import ( - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - is_sequence_parallel_enabled, -) -from lightx2v_train.runtime.sequence_parallel import all_gather_sequence -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .video_trainer import VideoDmdTrainer - - -@TRAINER_REGISTER("video_ar_dmd") -class VideoArDmdTrainer(VideoDmdTrainer): - trainer_name = "video_ar_dmd" - allowed_model_names = {"wan_t2v_ar", "wan_t2v_14b_ar", "wan_ti2v_5b", "wan_ti2v_5b_ar"} - supports_diversity_loss = False - supports_real_data_fake = False - - def __init__(self, config): - super().__init__(config) - self.num_frame_per_chunk = int(self.dmd_config["num_frame_per_chunk"] if "num_frame_per_chunk" in self.dmd_config else self.model_config.get("num_frame_per_chunk", 3)) - self.same_step_across_blocks = bool(self.dmd_config.get("same_step_across_blocks", True)) - self.context_noise = float(self.dmd_config.get("context_noise", 0.0)) - self.sequence_parallel_cache = bool(self.dmd_config["sp_cache"] if "sp_cache" in self.dmd_config else False) - - def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): - transformer = self.model.denoiser_module() - if not hasattr(transformer, "_forward_inference"): - raise RuntimeError("ar_dmd requires the causal Wan transformer.") - - if xt is None: - xt = self.sample_initial_latents(latent_shape) - batch_size, _, num_frames, _, _ = xt.shape - if num_frames % self.num_frame_per_chunk != 0: - raise ValueError(f"ar_dmd latent frames={num_frames} must be divisible by num_frame_per_chunk={self.num_frame_per_chunk}.") - - use_sp_cache = self._use_sequence_parallel_cache() - sp_size = get_sequence_parallel_world_size() if use_sp_cache else 1 - sp_rank = get_sequence_parallel_rank() if use_sp_cache else 0 - if use_sp_cache and self.num_frame_per_chunk % sp_size != 0: - raise ValueError(f"training.dmd.num_frame_per_chunk={self.num_frame_per_chunk} must be divisible by sequence_parallel.size={sp_size} when training.dmd.sp_cache=true.") - - self.model.transformer.train() - output_chunks = [] - context = self.model._condition_to_context_tensor(condition, batch_size=batch_size) - frame_seq_length = self._frame_seq_length(xt) - kv_cache, crossattn_cache = self._new_caches(batch_size, xt.dtype, xt.device, num_frames, frame_seq_length, sequence_parallel_cache=use_sp_cache) - num_blocks = num_frames // self.num_frame_per_chunk - exit_indices = self._sample_exit_indices(num_blocks, len(self.denoising_steps), xt.device) - - current_start_frame = 0 - for block_idx in range(num_blocks): - current_num_frames = self.num_frame_per_chunk - if use_sp_cache: - local_num_frames = current_num_frames // sp_size - local_start_frame = current_start_frame + sp_rank * local_num_frames - else: - local_num_frames = current_num_frames - local_start_frame = current_start_frame - latents = xt[:, :, local_start_frame : local_start_frame + local_num_frames] - exit_idx = int(exit_indices[0] if self.same_step_across_blocks else exit_indices[block_idx]) - - x0 = None - for step_idx, current_timestep in enumerate(self.denoising_steps): - timestep = torch.full( - (batch_size, local_num_frames), - float(current_timestep), - device=xt.device, - dtype=torch.float32, - ) - enable_step_grad = grad_enabled and step_idx == exit_idx - context_mgr = torch.enable_grad if enable_step_grad else torch.no_grad - with context_mgr(): - flow_pred = self._forward_causal_chunk( - self.model, - latents, - timestep, - context, - kv_cache, - crossattn_cache, - current_start=current_start_frame * frame_seq_length, - cache_start=current_start_frame * frame_seq_length, - local_frame_offset=local_start_frame, - balanced_sequence_parallel=use_sp_cache, - ) - x0 = self._flow_to_x0(latents, flow_pred, timestep) - - if step_idx == exit_idx: - break - - next_timestep = torch.full( - (batch_size, local_num_frames), - float(self.denoising_steps[step_idx + 1]), - device=xt.device, - dtype=torch.float32, - ) - with torch.no_grad(): - latents = self._add_noise_by_timestep(x0, torch.randn_like(x0), next_timestep) - - output_chunks.append(all_gather_sequence(x0, dim=2) if use_sp_cache else x0) - - cache_latents = x0.detach() - cache_timestep = torch.full( - (batch_size, local_num_frames), - self.context_noise, - device=xt.device, - dtype=torch.float32, - ) - if self.context_noise > 0: - cache_latents = self._add_noise_by_timestep(cache_latents, torch.randn_like(cache_latents), cache_timestep) - with torch.no_grad(): - self._forward_causal_chunk( - self.model, - cache_latents, - cache_timestep, - context, - kv_cache, - crossattn_cache, - current_start=current_start_frame * frame_seq_length, - cache_start=current_start_frame * frame_seq_length, - local_frame_offset=local_start_frame, - balanced_sequence_parallel=use_sp_cache, - ) - - current_start_frame += current_num_frames - - return ( - torch.cat(output_chunks, dim=2).to(dtype=self.running_dtype), - *self._denoised_timestep_window(exit_indices), - ) - - def _use_sequence_parallel_cache(self): - enabled = bool(self.sequence_parallel_cache and is_sequence_parallel_enabled()) - if enabled and bool(getattr(self.model.denoiser_module(), "defer_kv_cache_updates", False)): - raise ValueError("training.dmd.sp_cache=true does not support model.defer_kv_cache_updates=true. Set model.defer_kv_cache_updates=false or disable training.dmd.sp_cache.") - return enabled - - def _denoised_timestep_window(self, exit_indices): - if not self.same_step_across_blocks: - return None, None - - exit_idx = int(exit_indices[0]) - denoised_timestep_from = self._raw_timestep_from_warped_step(self.denoising_steps[exit_idx]) - if exit_idx == len(self.denoising_steps) - 1: - denoised_timestep_to = 0 - else: - denoised_timestep_to = self._raw_timestep_from_warped_step(self.denoising_steps[exit_idx + 1]) - return denoised_timestep_from, denoised_timestep_to - - def _sample_exit_indices(self, num_blocks, num_steps, device): - count = 1 if self.same_step_across_blocks else num_blocks - if dist.is_available() and dist.is_initialized(): - if dist.get_rank() == 0: - indices = torch.randint(0, num_steps, (count,), device=device) - else: - indices = torch.empty(count, dtype=torch.long, device=device) - dist.broadcast(indices, src=0) - return indices.tolist() - return torch.randint(0, num_steps, (count,), device=device).tolist() - - def _forward_causal_chunk( - self, - model, - latents, - timestep, - context, - kv_cache, - crossattn_cache, - current_start, - cache_start, - local_frame_offset=0, - balanced_sequence_parallel=False, - ): - transformer = model.denoiser_module() - seq_len = model._sequence_length(latents) - forward_context = model.transformer_forward_context() if hasattr(model, "transformer_forward_context") else torch.no_grad() - with forward_context: - return transformer( - latents, - t=timestep, - context=context, - seq_len=seq_len, - kv_cache=kv_cache, - crossattn_cache=crossattn_cache, - current_start=current_start, - cache_start=cache_start, - local_frame_offset=local_frame_offset, - balanced_sequence_parallel=balanced_sequence_parallel, - ) - - def _new_caches(self, batch_size, dtype, device, num_frames, frame_seq_length, sequence_parallel_cache=False): - transformer = self.model.denoiser_module() - num_layers = int(getattr(transformer, "num_layers", len(transformer.blocks))) - num_heads = int(transformer.num_heads) - head_dim = int(transformer.dim // transformer.num_heads) - kv_num_heads = num_heads - if sequence_parallel_cache: - sp_size = get_sequence_parallel_world_size() - if num_heads % sp_size != 0: - raise ValueError(f"transformer.num_heads={num_heads} must be divisible by sequence_parallel.size={sp_size}.") - kv_num_heads = num_heads // sp_size - local_attn_size = int(getattr(transformer, "local_attn_size", -1)) - kv_cache_size = num_frames * frame_seq_length if local_attn_size == -1 else local_attn_size * frame_seq_length - - kv_cache = [] - crossattn_cache = [] - for _ in range(num_layers): - kv_cache.append( - { - "k": torch.zeros((batch_size, kv_cache_size, kv_num_heads, head_dim), dtype=dtype, device=device), - "v": torch.zeros((batch_size, kv_cache_size, kv_num_heads, head_dim), dtype=dtype, device=device), - "global_end_index": torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": torch.tensor([0], dtype=torch.long, device=device), - } - ) - crossattn_cache.append( - { - "k": torch.zeros((batch_size, self.model.max_sequence_length, num_heads, head_dim), dtype=dtype, device=device), - "v": torch.zeros((batch_size, self.model.max_sequence_length, num_heads, head_dim), dtype=dtype, device=device), - "is_init": False, - } - ) - return kv_cache, crossattn_cache - - def _frame_seq_length(self, latent): - _, _, _, latent_height, latent_width = latent.shape - patch_t, patch_h, patch_w = self.model.patch_size - if patch_t != 1: - raise ValueError(f"ar_dmd expects temporal patch size 1, got {patch_t}.") - return latent_height * latent_width // (patch_h * patch_w) - - def _sigma_from_timestep(self, timestep, dtype): - timesteps = self.denoising_scheduler.timesteps.to(device=timestep.device, dtype=torch.float32) - sigmas = self.denoising_scheduler.sigmas.to(device=timestep.device, dtype=dtype) - flat_timestep = timestep.flatten().float() - index = torch.argmin((timesteps.unsqueeze(0) - flat_timestep.unsqueeze(1)).abs(), dim=1) - return sigmas[index].reshape(timestep.shape) - - def _expand_frame_sigma(self, sigma, ndim): - return sigma.reshape(sigma.shape[0], 1, sigma.shape[1], *([1] * (ndim - 3))) - - def _flow_to_x0(self, xt, flow_pred, timestep): - sigma = self._sigma_from_timestep(timestep, xt.dtype) - sigma = self._expand_frame_sigma(sigma, xt.ndim) - return (xt - sigma * flow_pred).to(dtype=xt.dtype) - - def _add_noise_by_timestep(self, x0, noise, timestep): - sigma = self._sigma_from_timestep(timestep, x0.dtype) - sigma = self._expand_frame_sigma(sigma, x0.ndim) - return ((1.0 - sigma) * x0 + sigma * noise).to(dtype=x0.dtype) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/video_trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/video_trainer.py deleted file mode 100644 index 48e0c5177..000000000 --- a/lightx2v_train/lightx2v_train/trainers/dmd/video_trainer.py +++ /dev/null @@ -1,826 +0,0 @@ -import copy -import os -from functools import partial - -import torch -import torch.distributed.checkpoint as dcp -import torch.nn.functional as F -from loguru import logger -from torch.distributed.checkpoint.state_dict import ( - StateDictOptions, - get_model_state_dict, - set_model_state_dict, -) - -from lightx2v_train.model_zoo import build_model -from lightx2v_train.runtime.distributed import ( - barrier, - get_world_size, - is_main_process, - reduce_mean, -) -from lightx2v_train.runtime.parallel import ( - apply_parallel, - set_parallel_gradient_sync, -) -from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value -from lightx2v_train.schedulers import DMDFlowMatchingScheduler -from lightx2v_train.schedulers.flow_matching import CausalForcingFlowMatchScheduler -from lightx2v_train.tricks import ( - DiversitySetupContext, - DiversityStepContext, - DiversityTrainerConstraints, - DiversityTrick, - RealDataFakeSetupContext, - RealDataFakeStepContext, - RealDataFakeTrick, -) -from lightx2v_train.utils.constants import ( - LINGBOT_VIDEO_NEGATIVE_PROMPT, - WAN_NEGATIVE_PROMPT, -) -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .config import VideoDmdConfig -from .trainer import DmdTrainer - - -@TRAINER_REGISTER("video_dmd") -class VideoDmdTrainer(DmdTrainer): - trainer_name = "video_dmd" - supports_cdm = False - allowed_model_names = {"wan_t2v", "wan_t2v_14b", "wan_ti2v_5b"} - default_negative_prompt = WAN_NEGATIVE_PROMPT - default_lora_target_modules = ("q", "k", "v", "o", "ffn.0", "ffn.2") - supports_diversity_loss = True - supports_real_data_fake = True - - def __init__(self, config): - super().__init__(config) - model_name = self.model_config["name"] - if model_name not in self.allowed_model_names: - allowed = ", ".join(repr(name) for name in sorted(self.allowed_model_names)) - raise ValueError(f"{self.trainer_name} trainer currently requires model.name in {{{allowed}}}.") - for role, lora_config in (("student", self.student_lora_config), ("fake", self.fake_lora_config)): - if lora_config is None or "target_modules" in lora_config: - continue - if self.default_lora_target_modules is None: - raise ValueError(f"training.{role}.lora.target_modules must be set for {self.trainer_name} training.") - lora_config["target_modules"] = list(self.default_lora_target_modules) - - parsed = VideoDmdConfig.from_mapping( - config, - dmd_config=self.dmd_config, - student_config=self.student_config, - model_config=self.model_config, - num_inference_steps=self.num_inference_steps, - ) - self.parsed_video_dmd_config = parsed - self.num_train_timestep = parsed.num_train_timestep - self.denoising_step_list = parsed.denoising_step_list - self.num_inference_steps = len(self.denoising_step_list) - self.warp_denoising_step = parsed.warp_denoising_step - self.min_step = parsed.min_step - self.max_step = parsed.max_step - self.score_timestep_shift = parsed.score_timestep_shift - self.ts_schedule = parsed.ts_schedule - self.ts_schedule_max = parsed.ts_schedule_max - self.min_score_timestep = parsed.min_score_timestep - self.student_checkpoint_path = parsed.student_checkpoint_path - self.student_checkpoint_strict = parsed.student_checkpoint_strict - - self.diversity_trick = DiversityTrick.from_mapping(self.dmd_config.get("div_loss", {})) - self.diversity_trick.validate_trainer( - DiversityTrainerConstraints( - supported=self.supports_diversity_loss, - rollout_steps=self.num_inference_steps, - ), - self.trainer_name, - ) - diversity_scheduler = ( - DMDFlowMatchingScheduler( - self.config, - self.dmd_config, - ) - if self.diversity_trick.enabled - else None - ) - self.diversity_trick.setup(DiversitySetupContext(scheduler=diversity_scheduler)) - self.real_data_fake_trick = RealDataFakeTrick.from_mappings( - { - "main": ( - "fake_real", - self.dmd_config.get("fake_real", {}), - ) - } - ) - if self.real_data_fake_trick.enabled and not self.supports_real_data_fake: - raise ValueError(f"{self.trainer_name} does not support training.dmd.fake_real.") - self.real_data_fake_trick.validate_schedule( - mode="standard", - num_train_timestep=self.num_train_timestep, - ) - if self.real_data_fake_trick.enabled: - train_data_name = self.config["data"]["train"].get("name") - if train_data_name not in { - "latent_dataset", - "video_dataset", - }: - raise ValueError("Real-data fake training requires data.train.name=latent_dataset or video_dataset.") - self.fake_real_train_type = self.fake_train_type - self.fake_real_lora_config = copy.deepcopy(self.fake_lora_config) - self.fake_real_optimizer_config = copy.deepcopy(self.fake_optimizer_config) - - def _prepare_cached_condition(self, condition): - if torch.is_tensor(condition): - prompt_embed = condition - elif isinstance(condition, dict) and "prompt_embed" in condition: - prompt_embed = condition["prompt_embed"] - else: - raise KeyError("WAN DMD cached condition expects conditioning.positive.prompt_embed.") - - prompt_embed = prompt_embed.to(device=self.model.device, dtype=self.running_dtype) - if prompt_embed.ndim == 2: - prompt_embed = prompt_embed.unsqueeze(0) - return {"prompt_embed": prompt_embed} - - def _encode_conditions(self, sample): - conditioning = sample["conditioning"] - positive = conditioning.get("positive") - if positive is None: - return super()._encode_conditions(sample) - - with torch.no_grad(): - condition = self._prepare_cached_condition(positive) - if self.guidance_scale > 1: - negative = conditioning.get("negative") - if negative is not None: - negative_condition = self._prepare_cached_condition(negative) - else: - batch_size = int(condition["prompt_embed"].shape[0]) - negative_prompt = self._negative_prompt_for_conditioning( - conditioning, - batch_size, - return_scalar=batch_size == 1, - ) - try: - negative_condition = self.model.encode_prompt_condition(negative_prompt) - except RuntimeError as exc: - raise RuntimeError( - "WAN DMD CFG needs negative conditions. Provide data.train.negative_condition_path or keep model.load_text_encoder=true for online negative prompt encoding." - ) from exc - else: - negative_condition = None - - condition = broadcast_sequence_parallel_value(condition) - negative_condition = broadcast_sequence_parallel_value(negative_condition) if negative_condition is not None else None - return condition, negative_condition - - def setup(self, resume_ckpt_path=None): - if resume_ckpt_path is None and self.student_checkpoint_path: - self._load_student_checkpoint(self.student_checkpoint_path, strict=self.student_checkpoint_strict) - super().setup(resume_ckpt_path=None) - self._setup_fake_real_resources() - if resume_ckpt_path is not None: - self._load_resume_state(resume_ckpt_path) - if resume_ckpt_path is None: - self.lr_scheduler = self._build_lr_scheduler( - self.optimizer, - num_training_steps=max(1, self.max_train_iters), - ) - self.fake_lr_scheduler = self._build_lr_scheduler( - self.fake_optimizer, - num_warmup_steps=0, - num_training_steps=max( - 1, - self.max_train_iters * self.fake_update_ratio, - ), - ) - - time_shift_settings = self.config["scheduler"].get("time_shift_settings", {}) - self.denoising_scheduler = CausalForcingFlowMatchScheduler( - num_train_timesteps=self.config["scheduler"].get("num_train_timesteps", 1000), - time_shift_settings=time_shift_settings, - ) - self.denoising_steps = self._build_denoising_steps(self.model.device) - self.denoising_sigmas = (self.denoising_steps / self.num_train_timestep).to(dtype=torch.float32) - self.real_data_fake_trick.setup(self._real_data_fake_setup_context()) - logger.info( - "[train] {} denoising_steps={} warped={}", - self.trainer_name, - [round(float(step), 4) for step in self.denoising_steps.detach().cpu()], - self.warp_denoising_step, - ) - - def _setup_fake_real_resources(self): - self.fake_real_model = None - self.fake_real_optimizer = None - self.fake_real_lr_scheduler = None - self.fake_real_trainable_params = [] - if not self.real_data_fake_trick.enabled_for("main"): - return - model_config = copy.deepcopy(self.fake_model_config) - self.fake_real_model = build_model(model_config) - self.fake_real_model.load_components( - transformer_only=True, - reference_model=self.model, - ) - self._setup_trainable_model( - self.fake_real_model, - role="fake_real", - ) - apply_parallel(self.fake_real_model, self.config) - if self.gradient_checkpointing: - self.fake_real_model.enable_gradient_checkpointing() - self.fake_real_trainable_params = list(self.fake_real_model.trainable_parameters()) - self.fake_real_optimizer = self._build_optimizer( - self.fake_real_trainable_params, - self.fake_real_optimizer_config, - ) - self.fake_real_lr_scheduler = self._build_lr_scheduler( - self.fake_real_optimizer, - num_warmup_steps=0, - num_training_steps=max( - 1, - self.max_train_iters * self.fake_update_ratio, - ), - ) - logger.info( - "[train] {} independent fake_real model={} path={} train_type={} trainable_params={}", - self.trainer_name, - model_config["model"]["name"], - model_config["model"]["pretrained_model_name_or_path"], - self.fake_real_train_type, - len(self.fake_real_trainable_params), - ) - - def _load_student_checkpoint(self, checkpoint_path, strict=True): - model_state_path = checkpoint_path - if os.path.isdir(model_state_path): - dist_state_path = os.path.join(model_state_path, "dist_state") - if os.path.isdir(dist_state_path): - module = self.model.fsdp2_state_module() - options = StateDictOptions(ignore_frozen_params=False, strict=strict) - model_state = get_model_state_dict(module, options=options) - state = {"model": model_state} - dcp.load(state, checkpoint_id=dist_state_path) - set_model_state_dict( - module, - model_state_dict=state["model"], - options=options, - ) - logger.info("[train] loaded {} student checkpoint from distributed state {}", self.trainer_name, dist_state_path) - return - model_state_path = os.path.join(model_state_path, "model_state.pt") - if not os.path.exists(model_state_path): - raise RuntimeError(f"{self.trainer_name} student checkpoint not found: {checkpoint_path}") - - state = torch.load(model_state_path, map_location="cpu", weights_only=False) - for key in ("generator_ema", "generator", "model", "state_dict"): - if isinstance(state, dict) and key in state: - state = state[key] - break - - fixed = {} - for key, value in state.items(): - for prefix in ( - "model._fsdp_wrapped_module.", - "model._checkpoint_wrapped_module.", - "model._orig_mod.", - "model.", - "_fsdp_wrapped_module.", - "_checkpoint_wrapped_module.", - "_orig_mod.", - ): - if key.startswith(prefix): - key = key[len(prefix) :] - fixed[key] = value - - incompatible = self.model.denoiser_module().load_state_dict(fixed, strict=strict) - if not strict: - if incompatible.missing_keys: - logger.warning("Missing keys when loading {} student checkpoint: {}", self.trainer_name, incompatible.missing_keys) - if incompatible.unexpected_keys: - logger.warning("Unexpected keys when loading {} student checkpoint: {}", self.trainer_name, incompatible.unexpected_keys) - logger.info("[train] loaded {} student checkpoint path={}", self.trainer_name, model_state_path) - - def train(self): - resume_ckpt_path, current_iter = self._resolve_resume() - self.setup(resume_ckpt_path=resume_ckpt_path) - if is_main_process(): - os.makedirs(self.output_train_dir, exist_ok=True) - barrier() - - max_train_iters = self.max_train_iters - grad_accum_iters = max(1, int(self.gradient_accumulation_iters)) - save_every_iters = self.save_every_iters - save_total_limit = self.save_total_limit - - logger.info( - "[train] start method={} student_train_type={} fake_train_type={} iter={}/{} world_size={} grad_accum={} train_log_every_iters={} fake_update_ratio={}", - self.training_config.get("method", self.trainer_name), - self.student_train_type, - self.fake_train_type, - current_iter, - max_train_iters, - get_world_size(), - grad_accum_iters, - self.train_log_every_iters, - self.fake_update_ratio, - ) - logger.info( - "[train] {} diversity enabled={} weight={} teacher_steps={} anchor_step={}", - self.trainer_name, - self.diversity_trick.enabled, - self.diversity_trick.config.weight, - self.diversity_trick.config.teacher_inference_steps, - self.diversity_trick.config.anchor_step, - ) - real_data_config = self.real_data_fake_trick.config.regions["main"] - logger.info( - "[train] {} real_data_fake enabled={} weight={} timesteps={}", - self.trainer_name, - real_data_config.enabled, - real_data_config.weight, - list(real_data_config.timestep_list), - ) - if self.infer_every_iters: - self.inferencer.set_data(self.dataloader_eval) - if current_iter == 0: - self.run_inference(current_iter) - - samples = self._iter_train_samples() - while current_iter < max_train_iters: - student_result = self._train_one_stage( - samples, - stage="student", - grad_accum_iters=grad_accum_iters, - ) - loss_dmd_value = student_result["dmd"] - loss_div_value = student_result["div_loss"] - loss_real_dmd_value = student_result["real_dmd"] - loss_fake_value = 0.0 - loss_fake_real_value = 0.0 - for _ in range(self.fake_update_ratio): - fake_result = self._train_one_stage( - samples, - stage="fake", - grad_accum_iters=grad_accum_iters, - ) - loss_fake_value += fake_result["loss"] - loss_fake_real_value += fake_result["fake_real"] - loss_fake_value /= self.fake_update_ratio - loss_fake_real_value /= self.fake_update_ratio - - current_iter += 1 - display_fake = reduce_mean(loss_fake_value) - display_dmd = reduce_mean(loss_dmd_value) if loss_dmd_value is not None else None - display_div = reduce_mean(loss_div_value) - display_real_dmd = reduce_mean(loss_real_dmd_value) - display_fake_real = reduce_mean(loss_fake_real_value) - current_lr = self.lr_scheduler.get_last_lr()[0] - if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: - dmd_text = "nan" if display_dmd is None else f"{display_dmd:.6f}" - logger.info( - "[train] iter={}/{} dmd={} div_loss={:.6f} real_dmd={:.6f} fake={:.6f} fake_real={:.6f} lr={:.8f}", - current_iter, - max_train_iters, - dmd_text, - display_div, - display_real_dmd, - display_fake, - display_fake_real, - current_lr, - ) - metrics = { - "train/div_loss": display_div, - "train/real_dmd": display_real_dmd, - "train/fake": display_fake, - "train/fake_real": display_fake_real, - "train/lr": current_lr, - } - if display_dmd is not None: - metrics["train/dmd"] = display_dmd - self.log_metrics(metrics, step=current_iter) - - if save_every_iters and current_iter % save_every_iters == 0: - self.save_checkpoint(current_iter, save_total_limit) - - if self.infer_every_iters and current_iter % self.infer_every_iters == 0: - self.run_inference(current_iter) - - logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) - - def _train_one_stage(self, samples, stage, grad_accum_iters): - if stage == "student": - optimizer = self.optimizer - scheduler = self.lr_scheduler - params = self.trainable_params - set_sync = self._set_student_gradient_sync - elif stage == "fake": - optimizer = self.fake_optimizer - scheduler = self.fake_lr_scheduler - params = self.fake_trainable_params - set_sync = self._set_fake_gradient_sync - else: - raise ValueError(f"Unsupported {self.trainer_name} training stage: {stage}") - - optimizer.zero_grad(set_to_none=True) - loss_value = 0.0 - use_real_data_fake = self.real_data_fake_trick.enabled_for("main") - if stage == "fake" and use_real_data_fake: - self.fake_real_optimizer.zero_grad(set_to_none=True) - if stage == "student": - metric_values = { - "div_loss": 0.0, - "real_dmd": 0.0, - } - else: - metric_values = {"fake_real": 0.0} - for micro_idx in range(grad_accum_iters): - sample = next(samples) - conditions = self._encode_conditions(sample) - latent_shape = self._latent_shape(sample) - initial_noise = self.sample_initial_latents(latent_shape) - set_sync((stage == "student" and use_real_data_fake) or micro_idx == grad_accum_iters - 1) - result = self.forward_loss( - latent_shape, - conditions, - stage=stage, - initial_noise=initial_noise, - ) - if isinstance(result, dict): - loss = result["loss"] - for name, value in result.items(): - if name == "loss": - continue - metric_values[name] = metric_values.get(name, 0.0) + value.item() / grad_accum_iters - del result - else: - loss = result - (loss / grad_accum_iters).backward() - loss_value += loss.item() / grad_accum_iters - del loss - if stage == "student" and self.diversity_trick.enabled: - div_raw, _ = self._backward_diversity_loss( - initial_noise, - conditions, - grad_accum_iters, - ) - metric_values["div_loss"] += div_raw - del initial_noise - if stage == "student" and use_real_data_fake: - real_result = self.real_data_fake_trick.student_loss( - self._real_data_fake_context( - sample, - conditions, - ) - ) - (real_result.loss / grad_accum_iters).backward() - metric_values["real_dmd"] += real_result.metrics["real_dmd"].item() / grad_accum_iters - del real_result - elif stage == "fake" and use_real_data_fake: - self._set_fake_real_gradient_sync(micro_idx == grad_accum_iters - 1) - fake_real_result = self.real_data_fake_trick.fake_loss( - self._real_data_fake_context( - sample, - conditions, - ) - ) - (fake_real_result.loss / grad_accum_iters).backward() - metric_values["fake_real"] += fake_real_result.metrics["fake_real"].item() / grad_accum_iters - del fake_real_result - - self._sync_sequence_parallel_grads(params) - torch.nn.utils.clip_grad_norm_(params, self.max_grad_norm) - optimizer.step() - if stage == "student": - self._after_student_optimizer_step("main") - scheduler.step() - optimizer.zero_grad(set_to_none=True) - if stage == "fake" and use_real_data_fake: - self._sync_sequence_parallel_grads(self.fake_real_trainable_params) - torch.nn.utils.clip_grad_norm_( - self.fake_real_trainable_params, - self.max_grad_norm, - ) - self.fake_real_optimizer.step() - self.fake_real_lr_scheduler.step() - self.fake_real_optimizer.zero_grad(set_to_none=True) - if metric_values: - return { - "loss": loss_value, - **metric_values, - } - return loss_value - - def _build_denoising_steps(self, device): - raw_steps = torch.tensor(self.denoising_step_list, dtype=torch.long, device=device) - if not self.warp_denoising_step: - return raw_steps.to(dtype=torch.float32) - - timesteps = torch.cat( - [ - self.denoising_scheduler.timesteps.to(device=device, dtype=torch.float32), - torch.zeros(1, device=device, dtype=torch.float32), - ] - ) - indices = self.denoising_scheduler.num_train_timesteps - raw_steps - return timesteps[indices] - - def _latent_shape(self, sample): - prompt = sample["conditioning"].get("prompt", "") - batch_size = 1 if isinstance(prompt, str) else len(prompt) - if "image_or_video_shape" in self.dmd_config: - configured_shape = self.dmd_config["image_or_video_shape"] - elif "image_or_video_shape" in self.training_config: - configured_shape = self.training_config["image_or_video_shape"] - else: - configured_shape = None - if configured_shape is not None: - shape = list(configured_shape) - shape[0] = batch_size - return tuple(int(dim) for dim in shape) - - missing_shape_keys = [key for key in ("height", "width", "num_frames") if key not in self.dmd_config] - if missing_shape_keys: - raise ValueError(f"training.dmd must define video shape fields: {', '.join(missing_shape_keys)}.") - height = int(self.dmd_config["height"]) - width = int(self.dmd_config["width"]) - num_frames = int(self.dmd_config["num_frames"]) - num_latent_frames = (int(num_frames) - 1) // self.model.vae_scale_factor_temporal + 1 - return ( - batch_size, - int(self.model._latent_channels()), - num_latent_frames, - int(height) // self.model.vae_scale_factor_spatial, - int(width) // self.model.vae_scale_factor_spatial, - ) - - def _extract_real_latents(self, sample): - with torch.no_grad(): - latent = sample["inputs"].get("latents") - if latent is None: - latent = self.model.encode_to_latent(sample) - latent = latent.to( - device=self.model.device, - dtype=self.running_dtype, - ) - if latent.ndim == 4: - latent = latent.unsqueeze(0) - return broadcast_sequence_parallel_value(latent) - - def _predict_real_student_velocity( - self, - region, - xt, - sigma, - condition, - grad_enabled, - ): - del region - context = torch.enable_grad if grad_enabled else torch.no_grad - self.model.transformer.train(mode=grad_enabled) - with context(): - return self._predict_velocity( - self.model, - xt, - sigma, - condition, - ) - - def _predict_real_teacher_velocity( - self, - region, - xt, - sigma, - condition, - negative_condition, - ): - del region - self.teacher_model.transformer.eval() - return self._predict_teacher_velocity( - xt, - sigma, - condition, - negative_condition, - ) - - def _diversity_step_context(self, initial_noise, conditions): - return DiversityStepContext( - initial_noise=initial_noise.detach(), - condition=conditions[0], - negative_condition=conditions[1], - latent_hw=initial_noise.shape[-2:], - device=self.model.device, - dtype=self.running_dtype, - predict_teacher_velocity=self._predict_teacher_velocity, - predict_student_velocity=partial( - self._predict_velocity, - self.model, - ), - student_scheduler=self.scheduler, - expand_to_ndim=self.scheduler._expand_to_ndim, - ) - - def _backward_diversity_loss( - self, - initial_noise, - conditions, - grad_accum_iters, - ): - self.diversity_trick.prepare_models( - self.teacher_model.transformer, - self.model.transformer, - ) - result = self.diversity_trick.student_loss(self._diversity_step_context(initial_noise, conditions)) - (result.loss / grad_accum_iters).backward() - raw_value = result.metrics["div_loss"].item() / grad_accum_iters - weighted_value = result.loss.detach().item() / grad_accum_iters - return raw_value, weighted_value - - def _real_data_fake_setup_context(self): - return RealDataFakeSetupContext( - mode="standard", - scheduler=self.scheduler, - denoising_scheduler=self.denoising_scheduler, - num_train_timestep=self.num_train_timestep, - warp_denoising_step=self.warp_denoising_step, - score_timestep_shift=self.score_timestep_shift, - min_step=self.min_step, - max_step=self.max_step, - min_score_timestep=self.min_score_timestep, - ) - - def _real_data_fake_context(self, sample, conditions, region="main"): - return RealDataFakeStepContext( - region=region, - fake_model=self.fake_real_model, - sample=sample, - condition=conditions[0], - negative_condition=conditions[1], - device=self.model.device, - dtype=self.running_dtype, - extract_real_latents=self._extract_real_latents, - sample_synced_int=self._sample_synced_int, - broadcast_noise=broadcast_sequence_parallel_value, - predict_student_velocity=self._predict_real_student_velocity, - predict_fake_velocity=self._predict_velocity, - predict_teacher_velocity=(self._predict_real_teacher_velocity), - dmd_loss=self._dmd_loss, - ) - - def _set_fake_real_gradient_sync(self, enabled): - if self.fake_real_model is not None: - set_parallel_gradient_sync( - self.fake_real_model, - enabled, - ) - - def _set_gradient_sync(self, enabled): - super()._set_gradient_sync(enabled) - self._set_fake_real_gradient_sync(enabled) - - def forward_loss( - self, - latent_shape, - conditions, - stage, - initial_noise=None, - ): - condition, negative_condition = conditions - ( - generated, - denoised_timestep_from, - denoised_timestep_to, - ) = self.run_back_simulation( - condition, - latent_shape, - grad_enabled=(stage != "fake"), - xt=initial_noise, - ) - - sigma = self._sample_score_sigma( - latent_shape[0], - denoised_timestep_from=denoised_timestep_from, - denoised_timestep_to=denoised_timestep_to, - device=self.model.device, - dtype=self.running_dtype, - ) - noise = torch.randn(latent_shape, device=self.model.device, dtype=torch.float32) - noise = broadcast_sequence_parallel_value(noise) - - with torch.no_grad(): - renoised_xt = self.scheduler.add_noise(generated, noise, sigma) - - if stage == "fake": - self.fake_model.transformer.train() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition) - velocity_gt = self.scheduler.build_train_gt(generated.float(), noise) - return F.mse_loss(velocity_fake.float(), velocity_gt.float(), reduction="mean") - - with torch.no_grad(): - self.fake_model.transformer.eval() - self.teacher_model.transformer.eval() - velocity_fake = self._predict_velocity(self.fake_model, renoised_xt, sigma, condition) - velocity_teacher_cond = self._predict_velocity(self.teacher_model, renoised_xt, sigma, condition) - if self.guidance_scale == 0: - velocity_teacher = velocity_teacher_cond - else: - velocity_teacher_uncond = self._predict_velocity(self.teacher_model, renoised_xt, sigma, negative_condition) - velocity_teacher = self._do_cfg(velocity_teacher_cond, velocity_teacher_uncond, self.guidance_scale, self.cfg_norm) - - expanded_sigma = self.scheduler._expand_to_ndim(sigma, renoised_xt.ndim) - x_pred_fake = renoised_xt - expanded_sigma * velocity_fake - x_pred_teacher = renoised_xt - expanded_sigma * velocity_teacher - - loss_dmd = self._dmd_loss( - generated, - x_pred_fake, - x_pred_teacher, - ) - return { - "loss": loss_dmd, - "dmd": loss_dmd.detach(), - } - - def sample_end_step(self): - return self._sample_synced_int( - self.diversity_trick.minimum_dmd_step_index, - self.scheduler.num_inference_steps, - ) - - def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): - transformer = self.model.denoiser_module() - if hasattr(transformer, "_forward_inference"): - raise RuntimeError("video_dmd requires the bidirectional Wan transformer. Use ar_dmd for causal student models.") - - self.scheduler.set_timesteps( - self.num_inference_steps, - sigmas=[float(sigma) for sigma in self.denoising_sigmas.detach().cpu()], - latent_hw=latent_shape[-2:], - device=self.model.device, - ) - if xt is None: - xt = self.sample_initial_latents(latent_shape) - - end_step_idx = ( - self.sample_end_step() - if grad_enabled - else self._sample_synced_int( - 0, - self.scheduler.num_inference_steps, - ) - ) - x0 = None - self.model.transformer.train() - for idx in range(end_step_idx + 1): - sigma = self.scheduler.sigma_at(idx, latent_shape[0], device=self.model.device, dtype=self.running_dtype) - context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad - with context(): - velocity = self._predict_velocity(self.model, xt, sigma, condition) - xt, x0 = self.scheduler.step_by_index(velocity, idx, xt) - - return x0.to(dtype=self.running_dtype), *self._denoised_timestep_window(end_step_idx) - - def _sample_score_sigma(self, batch_size, denoised_timestep_from, denoised_timestep_to, device, dtype): - min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep - max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep - min_timestep = max(0, int(min_timestep)) - max_timestep = min(self.num_train_timestep, int(max_timestep)) - if max_timestep <= min_timestep: - max_timestep = min(self.num_train_timestep, min_timestep + 1) - - timestep = torch.randint(min_timestep, max_timestep, (int(batch_size),), device=device, dtype=torch.long).float() - if self.score_timestep_shift > 1: - t = timestep / self.num_train_timestep - timestep = self.score_timestep_shift * t / (1 + (self.score_timestep_shift - 1) * t) * self.num_train_timestep - timestep = timestep.clamp(self.min_step, self.max_step) - return broadcast_sequence_parallel_value((timestep / self.num_train_timestep).to(dtype=dtype)) - - def _denoised_timestep_window(self, exit_idx): - exit_idx = int(exit_idx) - denoised_timestep_from = self._raw_timestep_from_warped_step(self.denoising_steps[exit_idx]) - if exit_idx == len(self.denoising_steps) - 1: - denoised_timestep_to = 0 - else: - denoised_timestep_to = self._raw_timestep_from_warped_step(self.denoising_steps[exit_idx + 1]) - return denoised_timestep_from, denoised_timestep_to - - def _raw_timestep_from_warped_step(self, warped_step): - if not self.warp_denoising_step: - return int(round(float(warped_step))) - timesteps = self.denoising_scheduler.timesteps.to(device=warped_step.device, dtype=torch.float32) - index = torch.argmin((timesteps - warped_step.float()).abs(), dim=0).item() - return self.denoising_scheduler.num_train_timesteps - int(index) - - -@TRAINER_REGISTER("lingbot_video_dmd") -class LingBotVideoDmdTrainer(VideoDmdTrainer): - trainer_name = "lingbot_video_dmd" - allowed_model_names = {"lingbot_video"} - default_negative_prompt = LINGBOT_VIDEO_NEGATIVE_PROMPT - default_lora_target_modules = None - - def _prepare_cached_condition(self, condition): - return self.model.prepare_text_condition(condition) diff --git a/lightx2v_train/lightx2v_train/trainers/dopsd.py b/lightx2v_train/lightx2v_train/trainers/dopsd.py index 3e746d85a..8e77ef070 100644 --- a/lightx2v_train/lightx2v_train/trainers/dopsd.py +++ b/lightx2v_train/lightx2v_train/trainers/dopsd.py @@ -7,16 +7,18 @@ import torch import torch.distributed.checkpoint as dcp -import torch.nn.functional as F from diffusers.optimization import get_scheduler from loguru import logger from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict from lightx2v_train.infer import build_inferencer from lightx2v_train.infer.dopsd_trajectory_viz import save_student_teacher_trajectory_grid +from lightx2v_train.model_capabilities import ( + DopsdCapability, + DopsdStepContext, +) from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints from lightx2v_train.runtime.distributed import barrier, get_rank, get_world_size, is_distributed, is_main_process, reduce_mean -from lightx2v_train.runtime.parallel import apply_parallel from lightx2v_train.utils.registry import TRAINER_REGISTER from lightx2v_train.utils.utils import get_running_dtype @@ -25,6 +27,11 @@ @TRAINER_REGISTER("dopsd") class DopsdTrainer(BaseTrainer): + required_capabilities = ( + *BaseTrainer.required_capabilities, + DopsdCapability, + ) + def __init__(self, config): super().__init__(config) self.running_dtype = get_running_dtype(self.model_config["running_dtype"]) @@ -78,6 +85,10 @@ def __init__(self, config): resume_config = self.config.get("resume", {}) self.auto_resume = resume_config.get("auto_resume", False) + def set_model(self, model): + super().set_model(model) + self.dopsd = model.capabilities.require(DopsdCapability) + def _resolve_ema_decay(self, current_iter): if self.ema_decay_warmup is not None and self.ema_decay_warmup_iters > 0 and current_iter <= self.ema_decay_warmup_iters: return float(self.ema_decay_warmup) @@ -92,28 +103,27 @@ def _step_loss_weights_for(self, num_steps): return weights[:num_steps] def setup(self, resume_ckpt_path=None): - self.model.add_dual_lora( + self.dopsd.configure_adapters( self.lora_rank, self.lora_alpha, self.lora_target_modules, - student_adapter=self.student_adapter, - teacher_adapter=self.teacher_adapter, - init_teacher_from_student=resume_ckpt_path is None, + self.student_adapter, + self.teacher_adapter, + resume_ckpt_path is None, ) - self.model.set_dual_lora_trainable(self.student_adapter, self.teacher_adapter) - apply_parallel(self.model, self.config) + self.parallel.apply(self.config) if self.gradient_checkpointing: - self.model.enable_gradient_checkpointing() + self.trainable_model.enable_gradient_checkpointing() if self.infer_every_iters: self.inferencer = build_inferencer(self.config) self.inferencer.set_model(self.model) - self.model.log_model_structure() + self.trainable_model.log_structure() - self.trainable_params = list(self.model.trainable_parameters()) + self.trainable_params = list(self.dopsd.parameters()) self.optimizer = torch.optim.AdamW( self.trainable_params, lr=self.optimizer_learning_rate, @@ -135,21 +145,24 @@ def _teacher_lora_checkpoint_path(self, resume_ckpt_path): return os.path.join(resume_ckpt_path, "teacher", "pytorch_lora_weights.safetensors") def _load_resume_state(self, resume_ckpt_path): - if self.model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp(): self._load_distributed_state(resume_ckpt_path) else: self._load_single_process_state(resume_ckpt_path) teacher_weights_path = self._teacher_lora_checkpoint_path(resume_ckpt_path) if os.path.exists(teacher_weights_path): - self.model.load_lora_weights_for_resume( + self.dopsd.load_adapter( resume_ckpt_path, - adapter_name=self.teacher_adapter, + self.teacher_adapter, weights_subdir="teacher", ) logger.info("Restored teacher EMA LoRA from {}", teacher_weights_path) else: - self.model.copy_lora_adapter_weights(self.student_adapter, self.teacher_adapter) + self.dopsd.copy_adapter( + self.student_adapter, + self.teacher_adapter, + ) logger.warning( "Teacher LoRA not found in checkpoint {}; initialized teacher from student", resume_ckpt_path, @@ -162,7 +175,10 @@ def _load_single_process_state(self, resume_ckpt_path): state = torch.load(training_state_path, map_location="cpu", weights_only=False) self._validate_checkpoint_metadata(state, training_state_path, resume_ckpt_path) - self.model.load_lora_weights_for_resume(resume_ckpt_path, adapter_name=self.student_adapter) + self.dopsd.load_adapter( + resume_ckpt_path, + self.student_adapter, + ) self.optimizer.load_state_dict(state["optimizer"]) self.lr_scheduler.load_state_dict(state["lr_scheduler"]) logger.info("Restored training state from {}", training_state_path) @@ -179,7 +195,7 @@ def _load_distributed_state(self, resume_ckpt_path): self._validate_checkpoint_metadata(trainer_state, trainer_state_path, resume_ckpt_path) options = StateDictOptions(ignore_frozen_params=True, strict=False) - state_module = self.model.fsdp2_state_module() + state_module = self.parallel.state_module() model_state, optim_state = get_state_dict(state_module, self.optimizer, options=options) state = {"model": model_state, "optimizer": optim_state} dcp.load(state, checkpoint_id=dist_state_path) @@ -243,83 +259,29 @@ def _validate_checkpoint_metadata(self, state, state_path, resume_ckpt_path): raise RuntimeError(f"Cannot resume checkpoint with iteration={checkpoint_iteration} in {state_path}, expected iteration={expected_iteration} from {resume_ckpt_path}") def compute_loss_on_sample(self, sample, collect_trajectory=False): - image = sample["inputs"].get("target_image") - if image is None: - raise ValueError("D-OPSD training requires inputs.target_image in each sample.") - - image = image.to(device=self.model.device, dtype=self.running_dtype) - bsz = image.shape[0] - height, width = image.shape[2], image.shape[3] - latent_hw = (height // 16, width // 16) - t_scale = float(self.noise_scheduler.num_train_timesteps) - - with torch.no_grad(): - student_condition = self.model.encode_condition(sample) - teacher_condition = self.model.encode_prompt_text(self._teacher_edit_prompts(sample["conditioning"]["prompt"])) - teacher_image_latents, teacher_image_latent_ids = self.model.prepare_reference_image_latents(image) - latents_begin, latent_ids = self.model.prepare_dopsd_initial_latents(height, width, bsz) - - self.noise_scheduler.set_timesteps(self.num_training_steps, latent_hw=latent_hw) - timesteps = self.noise_scheduler.infer_timesteps - - latents_student = latents_begin - total_loss = 0.0 - num_steps = len(timesteps) - step_loss_weights = self._step_loss_weights_for(num_steps) - weight_sum = 0.0 - student_x0_traj = [] - teacher_x0_traj = [] - - for back_step in range(num_steps): - t = timesteps[back_step].expand(bsz) / t_scale - t = t.to(device=self.model.device, dtype=self.running_dtype) - - if back_step < num_steps - 1: - next_t = timesteps[back_step + 1].expand(bsz) / t_scale - else: - next_t = torch.zeros_like(t) - next_t = next_t.to(device=self.model.device, dtype=self.running_dtype) - dt = next_t - t - - latents_student = latents_student.detach().requires_grad_(True) - - with torch.no_grad(): - v_pred_teacher = self.model.predict_velocity( - latents_student, - t, - teacher_condition, - latent_ids, - self.teacher_adapter, - teacher_image_latents=teacher_image_latents, - teacher_image_latent_ids=teacher_image_latent_ids, - ) - latents_teacher_cur = latents_student - x_0_teacher = latents_teacher_cur + (0 - t).reshape(bsz, 1, 1) * v_pred_teacher - # latents_teacher = latents_teacher_cur + v_pred_teacher * dt.reshape(bsz, 1, 1) - - v_pred_student = self.model.predict_velocity( - latents_student, - t, - student_condition, - latent_ids, - self.student_adapter, - ) - latents_student_cur = latents_student - x_0_student = latents_student_cur + (0 - t).reshape(bsz, 1, 1) * v_pred_student - latents_student = latents_student_cur + v_pred_student * dt.reshape(bsz, 1, 1) - - loss_dopsd = F.mse_loss(x_0_student, x_0_teacher.detach(), reduction="mean") - step_weight = step_loss_weights[back_step] - total_loss = total_loss + loss_dopsd * step_weight - weight_sum += step_weight - if collect_trajectory: - student_x0_traj.append(x_0_student.detach()) - teacher_x0_traj.append(x_0_teacher.detach()) - - avg_loss = total_loss / weight_sum + result = self.dopsd.compute_loss( + sample, + DopsdStepContext( + scheduler=self.noise_scheduler, + num_training_steps=self.num_training_steps, + running_dtype=self.running_dtype, + student_adapter=self.student_adapter, + teacher_adapter=self.teacher_adapter, + teacher_prompts=self._teacher_edit_prompts, + step_loss_weights=self._step_loss_weights_for, + collect_trajectory=collect_trajectory, + ), + ) if collect_trajectory: - return avg_loss, student_x0_traj, teacher_x0_traj, latent_ids, height, width - return avg_loss + return ( + result.loss, + result.student_trajectory, + result.teacher_trajectory, + result.latent_ids, + result.height, + result.width, + ) + return result.loss @torch.no_grad() def _save_training_trajectory(self, current_iter, student_x0_traj, teacher_x0_traj, latent_ids): @@ -334,9 +296,9 @@ def _save_training_trajectory(self, current_iter, student_x0_traj, teacher_x0_tr num_steps = len(student_x0_traj) for step_idx, (x_0_student, x_0_teacher) in enumerate(zip(student_x0_traj, teacher_x0_traj)): logger.info("[train] trajectory decode iter={} step={}/{} student", current_iter, step_idx + 1, num_steps) - student_step_images.extend(self.model.decode_packed_x0_to_images(x_0_student, latent_ids)) + student_step_images.extend(self.dopsd.decode_trajectory(x_0_student, latent_ids)) logger.info("[train] trajectory decode iter={} step={}/{} teacher", current_iter, step_idx + 1, num_steps) - teacher_step_images.extend(self.model.decode_packed_x0_to_images(x_0_teacher, latent_ids)) + teacher_step_images.extend(self.dopsd.decode_trajectory(x_0_teacher, latent_ids)) save_path = traj_dir / "student_teacher_x0_traj.png" save_student_teacher_trajectory_grid(student_step_images, teacher_step_images, save_path) @@ -419,7 +381,7 @@ def train(self): self.lr_scheduler.step() self.optimizer.zero_grad() current_ema_decay = self._resolve_ema_decay(current_iter + 1) - self.model.ema_update_lora_adapter( + self.dopsd.ema_update( self.student_adapter, self.teacher_adapter, current_ema_decay, @@ -451,14 +413,17 @@ def train(self): barrier() if is_main_process(): logger.info("[train] saving trajectory iter={} (decoding {} x0 pairs)...", current_iter, len(student_x0_traj)) - self.model.set_denoiser_eval() + self.dopsd.set_eval() self._save_training_trajectory( current_iter, student_x0_traj, teacher_x0_traj, latent_ids, ) - self.model.set_dual_lora_trainable(self.student_adapter, self.teacher_adapter) + self.dopsd.set_training( + self.student_adapter, + self.teacher_adapter, + ) barrier() if save_every_iters and current_iter % save_every_iters == 0: @@ -475,7 +440,7 @@ def train(self): logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) def _set_gradient_sync(self, enabled): - self.model.set_fsdp2_gradient_sync(enabled) + self.parallel.set_gradient_sync(enabled) @torch.no_grad() def _run_teacher_inference(self, current_iter, iter_output_dir): @@ -491,8 +456,8 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): teacher_output_dir = os.path.join(iter_output_dir, "teacher") os.makedirs(teacher_output_dir, exist_ok=True) - self.model.set_active_adapter(self.teacher_adapter) - self.model.set_denoiser_eval() + self.dopsd.set_active_adapter(self.teacher_adapter) + self.dopsd.set_eval() num_slots = (len(samples) + world_size - 1) // world_size if is_distributed() else len(samples) saved_count = 0 @@ -519,17 +484,24 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): ) continue - image = reference_image.unsqueeze(0).to(device=self.model.device, dtype=self.running_dtype) + image = reference_image.unsqueeze(0).to( + device=self.dopsd.device, + dtype=self.running_dtype, + ) height, width = image.shape[2], image.shape[3] latent_hw = (height // 16, width // 16) teacher_prompts = self._teacher_edit_prompts(record["prompt"]) - teacher_condition = self.model.encode_prompt_text(teacher_prompts) - teacher_image_latents, teacher_image_latent_ids = self.model.prepare_reference_image_latents(image) + teacher_condition = self.dopsd.encode_prompt(teacher_prompts) + teacher_image_latents, teacher_image_latent_ids = self.dopsd.prepare_reference(image) seed = base_seed + i - generator = torch.Generator(device=self.model.device).manual_seed(seed) - latents, latent_ids = self.model.prepare_dopsd_initial_latents(height, width, 1, generator=generator) + generator = torch.Generator(device=self.dopsd.device).manual_seed(seed) + latents, latent_ids = self.dopsd.initial_latents( + height, + width, + generator=generator, + ) self.noise_scheduler.set_timesteps(num_inference_steps, latent_hw=latent_hw) timesteps = self.noise_scheduler.infer_timesteps @@ -546,15 +518,18 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): ) for back_step in range(num_steps): t = timesteps[back_step].expand(1) / t_scale - t = t.to(device=self.model.device, dtype=self.running_dtype) + t = t.to(device=self.dopsd.device, dtype=self.running_dtype) if back_step < num_steps - 1: next_t = timesteps[back_step + 1].expand(1) / t_scale else: next_t = torch.zeros_like(t) - next_t = next_t.to(device=self.model.device, dtype=self.running_dtype) + next_t = next_t.to( + device=self.dopsd.device, + dtype=self.running_dtype, + ) dt = next_t - t - v_pred = self.model.predict_velocity( + v_pred = self.dopsd.predict_velocity( latents, t, teacher_condition, @@ -565,7 +540,7 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): ) latents = latents + v_pred * dt.reshape(1, 1, 1) - images = self.model.decode_packed_x0_to_images(latents, latent_ids) + images = self.dopsd.decode_trajectory(latents, latent_ids) save_path = Path(teacher_output_dir) / f"{i:05d}.png" images[0].save(save_path) saved_count += 1 @@ -573,8 +548,16 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): barrier() if is_distributed(): - saved_tensor = torch.tensor(saved_count, device=self.model.device, dtype=torch.int64) - skipped_tensor = torch.tensor(skipped_count, device=self.model.device, dtype=torch.int64) + saved_tensor = torch.tensor( + saved_count, + device=self.dopsd.device, + dtype=torch.int64, + ) + skipped_tensor = torch.tensor( + skipped_count, + device=self.dopsd.device, + dtype=torch.int64, + ) torch.distributed.all_reduce(saved_tensor, op=torch.distributed.ReduceOp.SUM) torch.distributed.all_reduce(skipped_tensor, op=torch.distributed.ReduceOp.SUM) saved_count = saved_tensor.item() @@ -590,7 +573,7 @@ def run_inference(self, current_iter): base_output_dir = self.infer_config.get("output_dir", "./output_infer") iter_output_dir = os.path.join(base_output_dir, f"iter-{current_iter:09d}") - self.model.set_active_adapter(self.student_adapter) + self.dopsd.set_active_adapter(self.student_adapter) self.inferencer.output_infer_dir = iter_output_dir os.makedirs(iter_output_dir, exist_ok=True) logger.info("[train] running student inference iter={} output_dir={}", current_iter, iter_output_dir) @@ -599,7 +582,10 @@ def run_inference(self, current_iter): self._run_teacher_inference(current_iter, iter_output_dir) logger.info("[train] finished inference iter={}", current_iter) - self.model.set_dual_lora_trainable(self.student_adapter, self.teacher_adapter) + self.dopsd.set_training( + self.student_adapter, + self.teacher_adapter, + ) def save_checkpoint(self, iteration, save_total_limit): if is_main_process(): @@ -611,16 +597,20 @@ def save_checkpoint(self, iteration, save_total_limit): os.makedirs(save_dir, exist_ok=True) barrier() - self.model.save_lora_weights(save_dir, adapter_name=self.student_adapter) + self.dopsd.save_adapter(save_dir, self.student_adapter) barrier() - self.model.save_lora_weights(save_dir, adapter_name=self.teacher_adapter, weights_subdir="teacher") + self.dopsd.save_adapter( + save_dir, + self.teacher_adapter, + weights_subdir="teacher", + ) barrier() config_path = self.config.get("config_path") if is_main_process() and config_path is not None: shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) - if self.model.is_fsdp2_wrapped(): + if self.parallel.is_fsdp(): self._save_distributed_state(save_dir, iteration) barrier() logger.info("[train] saved checkpoint iter={} path={} (student + teacher EMA LoRA)", iteration, save_dir) @@ -652,7 +642,11 @@ def _save_distributed_state(self, save_dir, iteration): barrier() options = StateDictOptions(ignore_frozen_params=True, strict=False) - model_state, optim_state = get_state_dict(self.model.fsdp2_state_module(), self.optimizer, options=options) + model_state, optim_state = get_state_dict( + self.parallel.state_module(), + self.optimizer, + options=options, + ) dcp.save( {"model": model_state, "optimizer": optim_state}, checkpoint_id=dist_state_path, diff --git a/lightx2v_train/lightx2v_train/trainers/flow.py b/lightx2v_train/lightx2v_train/trainers/flow.py deleted file mode 100644 index 2fddec934..000000000 --- a/lightx2v_train/lightx2v_train/trainers/flow.py +++ /dev/null @@ -1,293 +0,0 @@ -import os - -import torch -from loguru import logger -from torch import Tensor - -from lightx2v_train.model_zoo.native.ltx2 import ( - AudioLatentShape, - AudioPatchifier, - Modality, - SpatioTemporalScaleFactors, - VideoLatentPatchifier, - VideoLatentShape, - get_pixel_coords, -) -from lightx2v_train.runtime.distributed import barrier, get_world_size, is_main_process, reduce_mean -from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value, sync_sequence_parallel_gradients -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .base import BaseTrainer - - -@TRAINER_REGISTER("flow") -class FlowMatchingTrainer(BaseTrainer): - trainer_name = "flow" - - def compute_loss_on_sample(self, sample): - with torch.no_grad(): - latent = self.model.encode_to_latent(sample) - latent = broadcast_sequence_parallel_value(latent) - n = latent.shape[0] - noise = torch.randn_like(latent, dtype=self.running_dtype) - noise = broadcast_sequence_parallel_value(noise) - latent_hw = (latent.shape[-2], latent.shape[-1]) - timestep_or_sigma = self.noise_scheduler.sample_timestep_or_sigma(n, latent_hw=latent_hw) - timestep_or_sigma = broadcast_sequence_parallel_value(timestep_or_sigma) - noisy_latent = self.noise_scheduler.add_noise(latent, noise, timestep_or_sigma) - condition = self.model.encode_condition(sample) - condition = broadcast_sequence_parallel_value(condition) - - denoiser_input = self.model.prepare_denoiser_input(noisy_latent, condition=condition) - prediction = self.model.denoise(denoiser_input, timestep_or_sigma, condition) - prediction = self.model.postprocess_denoiser_output(prediction, denoiser_input) - - target = self.noise_scheduler.build_train_gt(latent, noise) - loss = torch.mean(((prediction.float() - target.float()) ** 2).reshape(target.shape[0], -1), dim=1) - return loss.mean() - - def train(self): - resume_ckpt_path, current_iter = self._resolve_resume() - # Objectives with a curriculum (for example CM's shrinking time - # interval) need the optimizer-step index, including after resume. - self.current_train_iteration = current_iter - self.setup(resume_ckpt_path=resume_ckpt_path) - if is_main_process(): - os.makedirs(self.output_train_dir, exist_ok=True) - barrier() - - max_train_iters = self.max_train_iters - grad_accum_iters = self.gradient_accumulation_iters - max_grad_norm = self.max_grad_norm - save_every_iters = self.save_every_iters - save_total_limit = self.save_total_limit - grad_accum_counter = 0 - running_loss = 0.0 - running_metrics = {} - - logger.info( - "[train] start method={} train_type={} iter={}/{} world_size={} grad_accum={} train_log_every_iters={}", - self.trainer_name, - self.train_type, - current_iter, - max_train_iters, - get_world_size(), - grad_accum_iters, - self.train_log_every_iters, - ) - if self.infer_every_iters: - self.inferencer.set_data(self.dataloader_eval) - if current_iter == 0: - self.run_inference(current_iter) - - epoch = 0 - while current_iter < max_train_iters: - sampler = getattr(self.dataloader_train, "sampler", None) - if hasattr(sampler, "set_epoch"): - sampler.set_epoch(epoch) - - for sample in self.dataloader_train: - sync_grad = (grad_accum_counter + 1) % grad_accum_iters == 0 - self._set_gradient_sync(sync_grad) - - loss_result = self.compute_loss_on_sample(sample) - if isinstance(loss_result, dict): - loss = loss_result["loss"] - sample_metrics = loss_result.get("metrics", {}) - else: - loss = loss_result - sample_metrics = {} - (loss / grad_accum_iters).backward() - running_loss += loss.item() / grad_accum_iters - for name, value in sample_metrics.items(): - scalar = value.detach().item() if torch.is_tensor(value) else float(value) - running_metrics[name] = running_metrics.get(name, 0.0) + scalar / grad_accum_iters - - grad_accum_counter += 1 - if grad_accum_counter % grad_accum_iters != 0: - continue - - self._after_backward() - torch.nn.utils.clip_grad_norm_(self.trainable_params, max_grad_norm) - self.optimizer.step() - self.lr_scheduler.step() - self.optimizer.zero_grad() - - current_iter += 1 - self.current_train_iteration = current_iter - display_loss = reduce_mean(running_loss) - current_lr = self.lr_scheduler.get_last_lr()[0] - if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: - display_metrics = {name: reduce_mean(value) for name, value in running_metrics.items()} - metric_text = " ".join(f"{name}={value:.6f}" for name, value in sorted(display_metrics.items())) - if metric_text: - logger.info("[train] iter={}/{} loss={:.6f} {} lr={:.8f}", current_iter, max_train_iters, display_loss, metric_text, current_lr) - else: - logger.info("[train] iter={}/{} loss={:.6f} lr={:.8f}", current_iter, max_train_iters, display_loss, current_lr) - logged_metrics = { - "train/loss": display_loss, - "train/lr": current_lr, - } - logged_metrics.update({f"train/{name}": value for name, value in display_metrics.items()}) - self.log_metrics(logged_metrics, step=current_iter) - running_loss = 0.0 - running_metrics = {} - - if save_every_iters and current_iter % save_every_iters == 0: - self.save_checkpoint(current_iter, save_total_limit) - - if self.infer_every_iters and current_iter % self.infer_every_iters == 0: - self.run_inference(current_iter) - - if current_iter >= max_train_iters: - break - - epoch += 1 - - logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) - - def _after_backward(self): - sync_sequence_parallel_gradients(self.trainable_params) - - -@TRAINER_REGISTER("ltx_t2av_flow") -class LTX2T2AVFlowTrainer(FlowMatchingTrainer): - trainer_name = "ltx_t2av_flow" - - def __init__(self, config): - super().__init__(config) - ltx2_config = self.training_config.get("ltx2", {}) - self.video_loss_weight = float(ltx2_config.get("video_loss_weight", 1.0)) - self.audio_loss_weight = float(ltx2_config.get("audio_loss_weight", 1.0)) - self.default_fps = float(ltx2_config.get("default_fps", 24.0)) - self.video_patchifier = VideoLatentPatchifier(patch_size=1) - self.audio_patchifier = AudioPatchifier(patch_size=1) - - def compute_loss_on_sample(self, sample): - inputs = sample["inputs"] - video_data = inputs.get("video_latents") - audio_data = inputs.get("audio_latents") - conditions = sample["conditioning"].get("positive") - if video_data is None or audio_data is None or conditions is None: - raise KeyError("ltx_t2av_flow expects inputs.video_latents, inputs.audio_latents and conditioning.positive.") - - video_latents = self._video_latents(video_data) - audio_latents = self._audio_latents(audio_data) - video_tokens = self.video_patchifier.patchify(video_latents) - audio_tokens = self.audio_patchifier.patchify(audio_latents) - - video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) - video_input, video_target, video_timesteps, video_loss_mask, video_sigmas = self._initialize_noisy_target(video_tokens) - audio_input, audio_target, audio_timesteps, audio_loss_mask, audio_sigmas = self._initialize_noisy_target(audio_tokens, sigmas=video_sigmas) - - video_positions = self._get_video_positions( - num_frames=int(video_data["num_frames"][0].item()) if "num_frames" in video_data else video_latents.shape[2], - height=int(video_data["height"][0].item()) if "height" in video_data else video_latents.shape[3], - width=int(video_data["width"][0].item()) if "width" in video_data else video_latents.shape[4], - batch_size=video_tokens.shape[0], - fps=float(video_data["fps"][0].item()) if "fps" in video_data else self.default_fps, - device=video_tokens.device, - ) - audio_positions = self._get_audio_positions( - num_time_steps=audio_tokens.shape[1], - batch_size=audio_tokens.shape[0], - device=audio_tokens.device, - ) - - video_modality = Modality( - enabled=True, - latent=video_input, - sigma=video_sigmas, - timesteps=video_timesteps, - positions=video_positions, - context=video_context, - context_mask=context_mask, - ) - audio_modality = Modality( - enabled=True, - latent=audio_input, - sigma=audio_sigmas, - timesteps=audio_timesteps, - positions=audio_positions, - context=audio_context, - context_mask=context_mask, - ) - - with self.model.transformer_forward_context(): - video_pred, audio_pred = self.model.denoiser_module()( - video=video_modality, - audio=audio_modality, - perturbations=None, - ) - - video_loss = self._masked_mse(video_pred, video_target, video_loss_mask) - audio_loss = self._masked_mse(audio_pred, audio_target, audio_loss_mask) - loss = self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss - return loss.mean() - - def _video_latents(self, video_data): - latents = video_data["latents"].to(device=self.model.device, dtype=self.running_dtype) - if latents.ndim == 4: - latents = latents.unsqueeze(0) - if latents.ndim != 5: - raise ValueError(f"LTX2 video latents must be [B,C,F,H,W] or [C,F,H,W], got shape={tuple(latents.shape)}") - return latents.contiguous() - - def _audio_latents(self, audio_data): - latents = audio_data["latents"].to(device=self.model.device, dtype=self.running_dtype) - if latents.ndim == 3: - latents = latents.unsqueeze(0) - if latents.ndim != 4: - raise ValueError(f"LTX2 audio latents must be [B,C,T,F] or [C,T,F], got shape={tuple(latents.shape)}") - return latents.contiguous() - - def _initialize_noisy_target(self, latents: Tensor, sigmas=None): - batch_size, seq_len, _ = latents.shape - if sigmas is None: - sigmas = self.noise_scheduler.sample_timestep_or_sigma(batch_size, seq_len=seq_len) - sigmas = sigmas.to(device=latents.device, dtype=latents.dtype) - noise = torch.randn_like(latents) - sigmas_expanded = sigmas.view(-1, 1, 1) - noisy_latents = (1 - sigmas_expanded) * latents + sigmas_expanded * noise - targets = noise - latents - timesteps = sigmas.view(-1, 1).expand(batch_size, seq_len).clone() - loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=latents.device) - return noisy_latents, targets, timesteps, loss_mask, sigmas - - def _get_video_positions(self, num_frames, height, width, batch_size, fps, device): - latent_coords = self.video_patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape( - frames=num_frames, - height=height, - width=width, - batch=batch_size, - channels=128, - ), - device=device, - ) - pixel_coords = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=SpatioTemporalScaleFactors.default(), - causal_fix=True, - ).float() - pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps - return pixel_coords - - def _get_audio_positions(self, num_time_steps, batch_size, device): - return self.audio_patchifier.get_patch_grid_bounds( - output_shape=AudioLatentShape( - frames=num_time_steps, - mel_bins=16, - batch=batch_size, - channels=8, - ), - device=device, - ) - - @staticmethod - def _masked_mse(pred, targets, loss_mask): - target_len = targets.shape[1] - pred = pred[:, -target_len:, :] - mask = loss_mask[:, -target_len:].unsqueeze(-1).float() - loss = (pred.float() - targets.float()).pow(2) * mask - return loss.mean(dim=(-2, -1)) / mask.mean(dim=(-2, -1)).clamp(min=1e-8) diff --git a/lightx2v_train/lightx2v_train/trainers/flow_matching.py b/lightx2v_train/lightx2v_train/trainers/flow_matching.py new file mode 100644 index 000000000..bd2a2d87d --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/flow_matching.py @@ -0,0 +1,136 @@ +import os + +import torch +from loguru import logger + +from lightx2v_train.model_capabilities import ( + FlowMatchingSFTCapability, + SFTStepContext, +) +from lightx2v_train.runtime.distributed import barrier, get_world_size, is_main_process, reduce_mean +from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value, sync_sequence_parallel_gradients +from lightx2v_train.utils.registry import TRAINER_REGISTER + +from .base import BaseTrainer + + +@TRAINER_REGISTER("flow_matching") +class FlowMatchingTrainer(BaseTrainer): + trainer_name = "flow_matching" + required_capabilities = ( + *BaseTrainer.required_capabilities, + FlowMatchingSFTCapability, + ) + + def set_model(self, model): + super().set_model(model) + self.sft = model.capabilities.require(FlowMatchingSFTCapability) + + def compute_loss_on_sample(self, sample): + return self.sft.compute_loss( + sample, + SFTStepContext( + noise_scheduler=self.noise_scheduler, + running_dtype=self.running_dtype, + broadcast=broadcast_sequence_parallel_value, + ), + ) + + def train(self): + resume_ckpt_path, current_iter = self._resolve_resume() + # Objectives with a curriculum (for example CM's shrinking time + # interval) need the optimizer-step index, including after resume. + self.current_train_iteration = current_iter + self.setup(resume_ckpt_path=resume_ckpt_path) + if is_main_process(): + os.makedirs(self.output_train_dir, exist_ok=True) + barrier() + + max_train_iters = self.max_train_iters + grad_accum_iters = self.gradient_accumulation_iters + max_grad_norm = self.max_grad_norm + save_every_iters = self.save_every_iters + save_total_limit = self.save_total_limit + grad_accum_counter = 0 + running_loss = 0.0 + running_metrics = {} + + logger.info( + "[train] start method={} train_type={} iter={}/{} world_size={} grad_accum={} train_log_every_iters={}", + self.trainer_name, + self.train_type, + current_iter, + max_train_iters, + get_world_size(), + grad_accum_iters, + self.train_log_every_iters, + ) + if self.infer_every_iters: + self.inferencer.set_data(self.dataloader_eval) + if current_iter == 0: + self.run_inference(current_iter) + + epoch = 0 + while current_iter < max_train_iters: + sampler = getattr(self.dataloader_train, "sampler", None) + if hasattr(sampler, "set_epoch"): + sampler.set_epoch(epoch) + + for sample in self.dataloader_train: + sync_grad = (grad_accum_counter + 1) % grad_accum_iters == 0 + self._set_gradient_sync(sync_grad) + + loss_result = self.compute_loss_on_sample(sample) + loss = loss_result.loss + sample_metrics = loss_result.metrics + (loss / grad_accum_iters).backward() + running_loss += loss.item() / grad_accum_iters + for name, value in sample_metrics.items(): + scalar = value.detach().item() if torch.is_tensor(value) else float(value) + running_metrics[name] = running_metrics.get(name, 0.0) + scalar / grad_accum_iters + + grad_accum_counter += 1 + if grad_accum_counter % grad_accum_iters != 0: + continue + + self._after_backward() + torch.nn.utils.clip_grad_norm_(self.trainable_params, max_grad_norm) + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad() + + current_iter += 1 + self.current_train_iteration = current_iter + display_loss = reduce_mean(running_loss) + current_lr = self.lr_scheduler.get_last_lr()[0] + if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: + display_metrics = {name: reduce_mean(value) for name, value in running_metrics.items()} + metric_text = " ".join(f"{name}={value:.6f}" for name, value in sorted(display_metrics.items())) + if metric_text: + logger.info("[train] iter={}/{} loss={:.6f} {} lr={:.8f}", current_iter, max_train_iters, display_loss, metric_text, current_lr) + else: + logger.info("[train] iter={}/{} loss={:.6f} lr={:.8f}", current_iter, max_train_iters, display_loss, current_lr) + logged_metrics = { + "train/loss": display_loss, + "train/lr": current_lr, + } + logged_metrics.update({f"train/{name}": value for name, value in display_metrics.items()}) + self.log_metrics(logged_metrics, step=current_iter) + running_loss = 0.0 + running_metrics = {} + + if save_every_iters and current_iter % save_every_iters == 0: + self.save_checkpoint(current_iter, save_total_limit) + + if self.infer_every_iters and current_iter % self.infer_every_iters == 0: + self.run_inference(current_iter) + + if current_iter >= max_train_iters: + break + + epoch += 1 + + logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) + + def _after_backward(self): + sync_sequence_parallel_gradients(self.trainable_params) diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/__init__.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/__init__.py index a3479624c..e72661603 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/__init__.py @@ -1,10 +1,7 @@ import importlib _LAZY_EXPORTS = { - "VideoPhasedDmdTrainer": ( - ".trainer", - "VideoPhasedDmdTrainer", - ), + "PhasedDmdTrainer": (".trainer", "PhasedDmdTrainer"), } diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/checkpoint.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/checkpoint.py index 684ae0c2d..2d22c45c7 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/checkpoint.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/checkpoint.py @@ -8,9 +8,7 @@ from loguru import logger from torch.distributed.checkpoint.state_dict import ( StateDictOptions, - get_model_state_dict, get_state_dict, - set_model_state_dict, set_state_dict, ) @@ -38,29 +36,11 @@ def _trainable_role_states(self): def _save_model_weights(self, model, save_dir, role="student"): train_type = self._role_train_type(role) - if train_type == "lora": - model.save_lora_weights(save_dir) - return - if is_main_process(): - torch.save( - model.denoiser_module().state_dict(), - os.path.join(save_dir, "model_state.pt"), - ) + self._checkpoint(model).save_weights(save_dir, train_type) def _load_model_weights(self, model, save_dir, role="student"): train_type = self._role_train_type(role) - if train_type == "lora": - model.load_lora_weights_for_resume(save_dir) - return - model_state_path = os.path.join(save_dir, "model_state.pt") - if not os.path.exists(model_state_path): - raise RuntimeError(f"model_state.pt not found in {save_dir}") - state_dict = torch.load( - model_state_path, - map_location="cpu", - weights_only=False, - ) - model.denoiser_module().load_state_dict(state_dict) + self._checkpoint(model).load_weights(save_dir, train_type) def _role_weights_dir(self, root_dir, role): directory_name = self.role_registry.weight_directory_name(role) @@ -71,26 +51,7 @@ def _role_weights_dir(self, root_dir, role): def _copy_fake_low_high_from_fake(self): if self.fake_low_high_model is None: return - source = self.fake_model.fsdp2_state_module() - target = self.fake_low_high_model.fsdp2_state_module() - if self.fake_model.is_fsdp2_wrapped(): - options = StateDictOptions( - ignore_frozen_params=False, - strict=True, - ) - source_state = get_model_state_dict( - source, - options=options, - ) - set_model_state_dict( - target, - model_state_dict=source_state, - options=options, - ) - del source_state - else: - target.load_state_dict(source.state_dict(), strict=True) - logger.warning("Checkpoint has no fake_low_high state; initialized it from the restored High Fake weights.") + self._copy_role_model("fake", "fake_low_high") def _fast_forward_fake_low_high_scheduler(self, iteration): if self.fake_low_high_lr_scheduler is None: @@ -181,7 +142,7 @@ def _validate_phased_checkpoint_metadata( def _load_resume_state(self, resume_ckpt_path): models = tuple(model for _, model in self._trainable_role_models()) - if any(model.is_fsdp2_wrapped() for model in models): + if any(self._parallel(model).is_fsdp() for model in models): self._load_distributed_state(resume_ckpt_path) return self._load_single_process_state(resume_ckpt_path) @@ -502,7 +463,7 @@ def _load_distributed_state(self, resume_ckpt_path): role, ) model_state, optimizer_state = get_state_dict( - model.fsdp2_state_module(), + self._parallel(model).state_module(), optimizer, options=options, ) @@ -516,7 +477,7 @@ def _load_distributed_state(self, resume_ckpt_path): process_group=checkpoint_group, ) set_state_dict( - model.fsdp2_state_module(), + self._parallel(model).state_module(), optimizer, model_state_dict=role_state["model"], optim_state_dict=role_state["optimizer"], @@ -533,7 +494,7 @@ def _load_distributed_state(self, resume_ckpt_path): for role in required_roles: model, optimizer = role_state_by_name[role] model_state, optimizer_state = get_state_dict( - model.fsdp2_state_module(), + self._parallel(model).state_module(), optimizer, options=options, ) @@ -547,7 +508,7 @@ def _load_distributed_state(self, resume_ckpt_path): for role in required_roles: model, optimizer = role_state_by_name[role] set_state_dict( - model.fsdp2_state_module(), + self._parallel(model).state_module(), optimizer, model_state_dict=state[f"{role}_model"], optim_state_dict=state[f"{role}_optimizer"], @@ -678,7 +639,7 @@ def save_checkpoint(self, iteration, save_total_limit): role_models = self._trainable_role_models() for role, model in role_models: weights_dir = self._role_weights_dir(save_dir, role) - save_weights = self._role_train_type(role) == "lora" or not model.is_fsdp2_wrapped() + save_weights = self._role_train_type(role) == "lora" or not self._parallel(model).is_fsdp() if save_weights and role != "student" and is_main_process(): os.makedirs(weights_dir, exist_ok=True) barrier() @@ -703,7 +664,7 @@ def save_checkpoint(self, iteration, save_total_limit): os.path.join(save_dir, "config.yaml"), ) - if any(model.is_fsdp2_wrapped() for _, model in role_models): + if any(self._parallel(model).is_fsdp() for _, model in role_models): self._save_distributed_state(save_dir, iteration) self._finalize_checkpoint( save_dir, @@ -799,7 +760,7 @@ def _save_distributed_state(self, save_dir, iteration): role, ) model_state, optimizer_state = get_state_dict( - model.fsdp2_state_module(), + self._parallel(model).state_module(), optimizer, options=options, ) diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py index 5df3f3ce0..81bf62d43 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py @@ -6,7 +6,7 @@ Optional, ) -from ..dmd.config import VideoDmdConfig +from ..dmd.config import DmdScheduleConfig def parse_role_lora_config( @@ -28,7 +28,7 @@ def parse_role_lora_config( @dataclass(frozen=True) -class PhasedDmdConfig(VideoDmdConfig): +class PhasedDmdConfig(DmdScheduleConfig): """Parsed configuration unique to dual-region phased DMD.""" phased: Dict[str, Any] @@ -66,7 +66,7 @@ def from_mapping( default_lora_target_modules, ): if dmd_config.get("cdm", {}).get("enabled", False): - raise ValueError("video_phased_dmd does not support training.dmd.cdm.") + raise ValueError("phased_dmd does not support training.dmd.cdm.") phased = dmd_config.get("phased", {}) if not isinstance(phased, dict): raise ValueError("training.dmd.phased must be a mapping.") @@ -79,7 +79,7 @@ def from_mapping( inference_steps = infer_config.get("denoising_step_list") if not isinstance(inference_steps, list): - raise ValueError("video_phased_dmd requires inference.denoising_step_list.") + raise ValueError("phased_dmd requires inference.denoising_step_list.") inference_step_count = int( infer_config.get( "num_inference_steps", @@ -102,7 +102,7 @@ def from_mapping( if "training_target" in phased: raise ValueError("training.dmd.phased.training_target is no longer supported. Phased DMD always trains High and Low regions together.") if max_train_iters < 2: - raise ValueError("video_phased_dmd requires training.max_train_iters >= 2 so both High and Low regions are trained.") + raise ValueError("phased_dmd requires training.max_train_iters >= 2 so both High and Low regions are trained.") margin = int(phased.get("score_timestep_margin", 20)) score_min = int(phased["score_timestep_min"]) diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/math.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/math.py index febca8c5f..54dbbaa4b 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/math.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/math.py @@ -30,7 +30,6 @@ def raw_timesteps_to_sigmas( def phase_sigma( - batch_size, *, match_timestep, device, @@ -40,7 +39,7 @@ def phase_sigma( num_train_timestep, ): raw_timestep = torch.full( - (int(batch_size),), + (1,), match_timestep, device=device, dtype=torch.long, @@ -56,7 +55,6 @@ def phase_sigma( def sample_score_sigma_range( - batch_size, raw_min, raw_max, *, @@ -93,7 +91,7 @@ def sample_score_sigma_range( candidate_indices = torch.randint( 0, candidate_sigmas.numel(), - (int(batch_size),), + (1,), device=device, dtype=torch.long, ) diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py index 23b60391d..d7a0eaad3 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py @@ -2,11 +2,12 @@ from loguru import logger -from lightx2v_train.model_zoo import build_model -from lightx2v_train.runtime.parallel import ( - apply_parallel, - set_parallel_gradient_sync, +from lightx2v_train.model_capabilities import ( + DistillationCapability, + ParallelCapability, + TrainableModelCapability, ) +from lightx2v_train.model_zoo import build_loaded_model from lightx2v_train.tricks import IdaModelPair from ..dmd.roles import ( @@ -96,22 +97,11 @@ def setup_trainable_model(self, model, role): lora_config = self.fake_real_low_lora_config else: raise ValueError(f"Unsupported phased model role: {role}") - if train_type == "lora": - model.add_lora( - lora_config["rank"], - lora_config["alpha"], - lora_config.get("target_modules"), - ) - model.set_lora_trainable() - return - model.set_full_trainable() + model.ensure_capabilities().require(TrainableModelCapability).configure(train_type, lora_config or {}) def restore_trainable_model(self, model, role): train_type = self.runtime(role).train_type - if train_type == "lora": - model.set_lora_trainable() - return - model.set_full_trainable() + model.ensure_capabilities().require(TrainableModelCapability).restore(train_type) def role_model_config(self, role): excluded_roles = { @@ -141,12 +131,12 @@ def role_model_config(self, role): def ida_model_pairs(self): return { "high": IdaModelPair( - student=self.model.denoiser_module(), - fake=self.fake_model.denoiser_module(), + student=self.student.denoiser(), + fake=self.fake.denoiser(), ), "low": IdaModelPair( - student=self.student_2_model.denoiser_module(), - fake=self.fake_2_model.denoiser_module(), + student=self.student_2.denoiser(), + fake=self.fake_2.denoiser(), ), } @@ -199,7 +189,7 @@ def active_fake_real_state(self, region): def set_role_gradient_sync(self, role, enabled): runtime = self.runtime(role) - set_parallel_gradient_sync(runtime.model, enabled) + runtime.model.capabilities.require(ParallelCapability).set_gradient_sync(enabled) def trainable_names(self): names = ["student", "fake", "student_2", "fake_2"] @@ -253,48 +243,51 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): self.student_checkpoint_path = student_checkpoint_path student_2_model_config = self._role_model_config("student_2") - self.student_2_model = build_model(student_2_model_config) - self.student_2_model.load_components( + self.student_2_model = build_loaded_model( + student_2_model_config, transformer_only=True, reference_model=self.model, ) + self.student_2 = self.student_2_model.capabilities.require(DistillationCapability) self._setup_trainable_model( self.student_2_model, role="student_2", ) - apply_parallel(self.student_2_model, self.config) + self.student_2_model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: - self.student_2_model.enable_gradient_checkpointing() + self.student_2_model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() fake_2_model_config = self._role_model_config("fake_2") - self.fake_2_model = build_model(fake_2_model_config) - self.fake_2_model.load_components( + self.fake_2_model = build_loaded_model( + fake_2_model_config, transformer_only=True, reference_model=self.model, ) + self.fake_2 = self.fake_2_model.capabilities.require(DistillationCapability) self._setup_trainable_model(self.fake_2_model, role="fake_2") - apply_parallel(self.fake_2_model, self.config) + self.fake_2_model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: - self.fake_2_model.enable_gradient_checkpointing() + self.fake_2_model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() self.fake_low_high_model = None fake_low_high_model_config = None fake_low_high_model_path = None if self.enable_fake_low_high: fake_low_high_model_config = self._role_model_config("teacher") - self.fake_low_high_model = build_model(fake_low_high_model_config) - fake_low_high_model_path = fake_low_high_model_config["model"]["pretrained_model_name_or_path"] - self.fake_low_high_model.load_components( + self.fake_low_high_model = build_loaded_model( + fake_low_high_model_config, transformer_only=True, reference_model=self.model, ) + fake_low_high_model_path = fake_low_high_model_config["model"]["pretrained_model_name_or_path"] + self.fake_low_high = self.fake_low_high_model.capabilities.require(DistillationCapability) self._setup_trainable_model( self.fake_low_high_model, role="fake_low_high", ) - apply_parallel(self.fake_low_high_model, self.config) + self.fake_low_high_model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: - self.fake_low_high_model.enable_gradient_checkpointing() + self.fake_low_high_model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() self.fake_real_high_model = None self.fake_real_low_model = None @@ -305,16 +298,18 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): if not self.real_data_fake_trick.enabled_for(region): continue role_config = self._role_model_config(source_role) - model = build_model(role_config) - model.load_components( + model = build_loaded_model( + role_config, transformer_only=True, reference_model=self.model, ) self._setup_trainable_model(model, role=role) - apply_parallel(model, self.config) + model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: - model.enable_gradient_checkpointing() + model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() setattr(self, f"{role}_model", model) + capability = model.capabilities.require(DistillationCapability) + setattr(self, role, capability) logger.info( "[train] phased_dmd independent {} path={} train_type={}", role, @@ -332,24 +327,26 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): if teacher_model_config != teacher_2_model_config: raise ValueError("model.teacher_2 can share with teacher only when their model configurations match.") self.teacher_2_model = self.teacher_model + self.teacher_2 = self.teacher logger.info("[train] phased_dmd teacher_2 shares teacher weights") else: teacher_2_model_config = self._role_model_config("teacher_2") - self.teacher_2_model = build_model(teacher_2_model_config) - self.teacher_2_model.load_components( + self.teacher_2_model = build_loaded_model( + teacher_2_model_config, transformer_only=True, reference_model=self.model, ) - self.teacher_2_model.transformer.requires_grad_(False) - self.teacher_2_model.transformer.eval() - apply_parallel(self.teacher_2_model, self.config) - self.teacher_2_model.transformer.eval() - - self.student_2_trainable_params = list(self.student_2_model.trainable_parameters()) - self.fake_2_trainable_params = list(self.fake_2_model.trainable_parameters()) - self.fake_low_high_trainable_params = list(self.fake_low_high_model.trainable_parameters()) if self.fake_low_high_model is not None else [] - self.fake_real_high_trainable_params = list(self.fake_real_high_model.trainable_parameters()) if self.fake_real_high_model is not None else [] - self.fake_real_low_trainable_params = list(self.fake_real_low_model.trainable_parameters()) if self.fake_real_low_model is not None else [] + self.teacher_2 = self.teacher_2_model.capabilities.require(DistillationCapability) + self.teacher_2.denoiser().requires_grad_(False) + self.teacher_2.set_training(False) + self.teacher_2_model.capabilities.require(ParallelCapability).apply(self.config) + self.teacher_2.set_training(False) + + self.student_2_trainable_params = list(self.student_2_model.capabilities.require(TrainableModelCapability).parameters()) + self.fake_2_trainable_params = list(self.fake_2_model.capabilities.require(TrainableModelCapability).parameters()) + self.fake_low_high_trainable_params = list(self.fake_low_high_model.capabilities.require(TrainableModelCapability).parameters()) if self.fake_low_high_model is not None else [] + self.fake_real_high_trainable_params = list(self.fake_real_high_model.capabilities.require(TrainableModelCapability).parameters()) if self.fake_real_high_model is not None else [] + self.fake_real_low_trainable_params = list(self.fake_real_low_model.capabilities.require(TrainableModelCapability).parameters()) if self.fake_real_low_model is not None else [] self.student_2_optimizer = self._build_optimizer( self.student_2_trainable_params, self.student_2_optimizer_config, @@ -440,7 +437,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): ) if self.infer_every_iters: if not hasattr(self.inferencer, "set_low_model"): - raise RuntimeError("video_phased_dmd requires wan_t2v_dual_infer.") + raise RuntimeError("phased_dmd inference requires an inferencer that supports set_low_model().") self.inferencer.set_low_model(self.student_2_model) if resume_ckpt_path is not None: diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py index dc3ab2d90..7f0ee24d2 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py @@ -35,7 +35,7 @@ def _run_region_rollout( self.num_inference_steps, sigmas=[float(sigma) for sigma in self.denoising_sigmas.detach().cpu()], latent_hw=latent_shape[-2:], - device=self.model.device, + device=self.student.device, ) if region == "high": min_step_index = self.diversity_trick.minimum_dmd_step_index if grad_enabled else 0 @@ -51,14 +51,13 @@ def _run_region_rollout( if xt is None: xt = self.sample_initial_latents(latent_shape) - self.model.transformer.train() - self.student_2_model.transformer.train() + self.student.set_training(True) + self.student_2.set_training(True) for step_index in range(gradient_step_index + 1): - active_model = self.model if step_index < self.match_step_index else self.student_2_model + active_model = self.student if step_index < self.match_step_index else self.student_2 sigma = self.scheduler.sigma_at( step_index, - latent_shape[0], - device=self.model.device, + device=self.student.device, dtype=self.running_dtype, ) keep_gradient = grad_enabled and step_index == gradient_step_index @@ -86,7 +85,6 @@ def _run_region_rollout( ) if region == "high": sigma_s = self._phase_sigma( - latent_shape[0], dtype=self.running_dtype, ) anchor = self._euler_step( @@ -117,19 +115,18 @@ def _run_full_fake_rollout( self.num_inference_steps, sigmas=[float(sigma) for sigma in self.denoising_sigmas.detach().cpu()], latent_hw=latent_shape[-2:], - device=self.model.device, + device=self.student.device, ) if xt is None: xt = self.sample_initial_latents(latent_shape) - self.model.transformer.eval() - self.student_2_model.transformer.eval() + self.student.set_training(False) + self.student_2.set_training(False) x_bound = None for step_index in range(self.num_inference_steps): - active_model = self.model if step_index < self.match_step_index else self.student_2_model + active_model = self.student if step_index < self.match_step_index else self.student_2 sigma = self.scheduler.sigma_at( step_index, - latent_shape[0], - device=self.model.device, + device=self.student.device, dtype=self.running_dtype, ) velocity = self._predict_velocity( @@ -194,13 +191,12 @@ def _dmd_loss_for_region( if teacher_branch == "high": raw_min = self.match_timestep + self.score_timestep_margin raw_max = self.num_train_timestep - teacher_model = self.teacher_model + teacher_model = self.teacher else: raw_min = 1 raw_max = self.match_timestep - teacher_model = self.teacher_2_model + teacher_model = self.teacher_2 sigma_t = self._sample_score_sigma_range( - anchor.shape[0], raw_min, raw_max, anchor.device, @@ -218,8 +214,8 @@ def _dmd_loss_for_region( sigma_s, sigma_t, ) - fake_model.transformer.eval() - teacher_model.transformer.eval() + fake_model.set_training(False) + teacher_model.set_training(False) velocity_fake = self._predict_velocity( fake_model, score_xt, @@ -255,10 +251,10 @@ def _dmd_loss_for_region( def _fake_model_for_dmd(self, region, teacher_branch): if region == "high": - return self.fake_model + return self.fake if teacher_branch == "high" and self.fake_low_high_model is not None: - return self.fake_low_high_model - return self.fake_2_model + return self.fake_low_high + return self.fake_2 def _fake_loss_for_score_range( self, @@ -270,7 +266,6 @@ def _fake_loss_for_score_range( raw_max, ): sigma_t = self._sample_score_sigma_range( - anchor.shape[0], raw_min, raw_max, anchor.device, @@ -290,7 +285,7 @@ def _fake_loss_for_score_range( sigma_s, sigma_t, ) - fake_model.transformer.train() + fake_model.set_training(True) velocity_fake = self._predict_velocity( fake_model, score_xt, @@ -307,20 +302,19 @@ def _fake_score_specifications(self, region, x_bound, x0): if region == "high": anchor = x_bound sigma_s = self._phase_sigma( - x_bound.shape[0], dtype=self.running_dtype, ) - fake_model = self.fake_model + fake_model = self.fake raw_min = self.match_timestep + self.score_timestep_margin raw_max = self.num_train_timestep else: anchor = x0 sigma_s = torch.zeros( - x0.shape[0], + 1, device=x0.device, dtype=self.running_dtype, ) - fake_model = self.fake_2_model + fake_model = self.fake_2 raw_min = 1 raw_max = self.num_train_timestep if self.fake_low_high_model is None else self.match_timestep specifications = [ @@ -337,10 +331,10 @@ def _fake_score_specifications(self, region, x_bound, x0): specifications.append( ( "fake_low_high", - self.fake_low_high_model, + self.fake_low_high, x0, torch.zeros( - x0.shape[0], + 1, device=x0.device, dtype=self.running_dtype, ), @@ -351,18 +345,11 @@ def _fake_score_specifications(self, region, x_bound, x0): return tuple(specifications) def _extract_real_latents(self, sample): - with torch.no_grad(): - latent = sample["inputs"].get("latents") - if latent is None: - latent = self.model.encode_to_latent(sample) - latent = latent.to( - device=self.model.device, - dtype=self.running_dtype, - ) - if latent.ndim == 4: - latent = latent.unsqueeze(0) - latent = broadcast_sequence_parallel_value(latent) - return latent + return self.student.extract_real_latents( + sample, + self.running_dtype, + broadcast_sequence_parallel_value, + ) def _predict_real_student_velocity( self, @@ -374,8 +361,8 @@ def _predict_real_student_velocity( ): context = torch.enable_grad if grad_enabled else torch.no_grad with context(): - student_model = self.model if region == "high" else self.student_2_model - student_model.transformer.train(mode=grad_enabled) + student_model = self.student if region == "high" else self.student_2 + student_model.set_training(grad_enabled) velocity = self._predict_velocity( student_model, xt, @@ -392,8 +379,8 @@ def _predict_real_teacher_velocity( condition, negative_condition, ): - teacher_model = self.teacher_model if region == "high" else self.teacher_2_model - teacher_model.transformer.eval() + teacher_model = self.teacher if region == "high" else self.teacher_2 + teacher_model.set_training(False) return self._teacher_velocity( teacher_model, xt, @@ -408,15 +395,15 @@ def _diversity_step_context(self, initial_noise, conditions): condition=conditions[0], negative_condition=conditions[1], latent_hw=initial_noise.shape[-2:], - device=self.model.device, + device=self.student.device, dtype=self.running_dtype, predict_teacher_velocity=partial( self._teacher_velocity, - self.teacher_model, + self.teacher, ), predict_student_velocity=partial( self._predict_velocity, - self.model, + self.student, ), student_scheduler=self.scheduler, expand_to_ndim=self.scheduler._expand_to_ndim, @@ -439,11 +426,11 @@ def _real_data_fake_setup_context(self): def _real_data_fake_context(self, sample, conditions, region): return RealDataFakeStepContext( region=region, - fake_model=(self.fake_real_high_model if region == "high" else self.fake_real_low_model), + fake_model=(self.fake_real_high if region == "high" else self.fake_real_low), sample=sample, condition=conditions[0], negative_condition=conditions[1], - device=self.model.device, + device=self.student.device, dtype=self.running_dtype, extract_real_latents=self._extract_real_latents, sample_synced_int=self._sample_synced_int, diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/trainer.py index 9cd580281..d1fb8867b 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/trainer.py @@ -17,7 +17,7 @@ ) from lightx2v_train.utils.registry import TRAINER_REGISTER -from ..dmd import VideoDmdTrainer +from ..dmd import DmdTrainer from ..dmd.math import euler_step from .checkpoint import PhasedCheckpointManager from .config import ( @@ -36,11 +36,11 @@ from .rollout import PhasedRolloutEngine -@TRAINER_REGISTER("video_phased_dmd") -class VideoPhasedDmdTrainer(VideoDmdTrainer): +@TRAINER_REGISTER("phased_dmd") +class PhasedDmdTrainer(DmdTrainer): """Dual-student phased DMD with x0-space distribution matching.""" - trainer_name = "video_phased_dmd" + trainer_name = "phased_dmd" defer_ida_setup = True def __init__(self, config): @@ -50,11 +50,11 @@ def __init__(self, config): self.rollout_engine = PhasedRolloutEngine(self) parsed = PhasedDmdConfig.from_mapping( config, - base_config=self.parsed_video_dmd_config, + base_config=self.parsed_dmd_schedule_config, dmd_config=self.dmd_config, infer_config=self.infer_config, max_train_iters=self.max_train_iters, - default_lora_target_modules=(self.default_lora_target_modules), + default_lora_target_modules=None, ) self.parsed_phased_dmd_config = parsed self._checkpoint_process_group = None @@ -165,7 +165,7 @@ def setup(self, resume_ckpt_path=None): return self.role_registry.setup_resources( resume_ckpt_path=resume_ckpt_path, base_setup=super( - VideoPhasedDmdTrainer, + PhasedDmdTrainer, self, ).setup, ) @@ -188,11 +188,10 @@ def _teacher_branch_for_iteration(self, current_iter, region): low_iteration = current_iter // 2 return "high" if low_iteration % 2 == 0 else "low" - def _phase_sigma(self, batch_size, dtype): + def _phase_sigma(self, dtype): return phase_sigma( - batch_size, match_timestep=self.match_timestep, - device=self.model.device, + device=self.student.device, dtype=dtype, warp_denoising_step=self.warp_denoising_step, denoising_scheduler=self.denoising_scheduler, @@ -202,7 +201,7 @@ def _phase_sigma(self, batch_size, dtype): def _raw_timesteps_to_sigmas(self, raw_timesteps, dtype): return raw_timesteps_to_sigmas( raw_timesteps, - device=self.model.device, + device=self.student.device, dtype=dtype, warp_denoising_step=self.warp_denoising_step, denoising_scheduler=self.denoising_scheduler, @@ -211,14 +210,12 @@ def _raw_timesteps_to_sigmas(self, raw_timesteps, dtype): def _sample_score_sigma_range( self, - batch_size, raw_min, raw_max, device, dtype, ): return sample_score_sigma_range( - batch_size, raw_min, raw_max, device=device, diff --git a/lightx2v_train/lightx2v_train/trainers/sgmd.py b/lightx2v_train/lightx2v_train/trainers/sgmd.py index f4c525c55..4a9ff8ba6 100644 --- a/lightx2v_train/lightx2v_train/trainers/sgmd.py +++ b/lightx2v_train/lightx2v_train/trainers/sgmd.py @@ -17,7 +17,7 @@ ) from lightx2v_train.utils.registry import TRAINER_REGISTER -from .dmd import VideoDmdTrainer +from .dmd import DmdTrainer def normalized_fisher_loss( @@ -56,11 +56,11 @@ def generator_fake_correction_loss( ) -@TRAINER_REGISTER("video_sgmd") -class VideoSgmdTrainer(VideoDmdTrainer): +@TRAINER_REGISTER("sgmd") +class SgmdTrainer(DmdTrainer): """Strict video SGMD with one student and one fake-score update per iter.""" - trainer_name = "video_sgmd" + trainer_name = "sgmd" fake_correction_weight = 0.1 supports_real_data_fake = False supports_ida = False @@ -68,9 +68,9 @@ class VideoSgmdTrainer(VideoDmdTrainer): def __init__(self, config): super().__init__(config) if self.fake_update_ratio != 1: - raise ValueError("video_sgmd strictly uses one student and one fake-score update per iteration; set training.dmd.fake_update_ratio=1.") + raise ValueError("SGMD uses one student and one fake-score update per iteration; set training.dmd.fake_update_ratio=1.") if self.cdm_enabled: - raise ValueError("video_sgmd does not support training.dmd.cdm.") + raise ValueError("SGMD does not support training.dmd.cdm.") def sample_end_step(self): end_step_idx = self._sample_synced_int( @@ -99,16 +99,15 @@ def _generator_forward( ) sgmd_step = self._last_sgmd_step sigma = self._sample_score_sigma( - latent_shape[0], denoised_timestep_from=denoised_timestep_from, denoised_timestep_to=denoised_timestep_to, - device=self.model.device, + device=self.student.device, dtype=self.running_dtype, ) noise = broadcast_sequence_parallel_value( torch.randn( latent_shape, - device=self.model.device, + device=self.student.device, dtype=torch.float32, ) ) @@ -118,15 +117,15 @@ def _generator_forward( sigma, ) - self.fake_model.transformer.eval() + self.fake.set_training(False) velocity_fake = self._predict_velocity( - self.fake_model, + self.fake, renoised_xt, sigma, condition, ) with torch.no_grad(): - self.teacher_model.transformer.eval() + self.teacher.set_training(False) velocity_teacher = self._predict_teacher_velocity( renoised_xt, sigma, @@ -174,9 +173,9 @@ def _fake_score_forward(self, score_context): noise, sigma, ) - self.fake_model.transformer.eval() + self.fake.set_training(False) velocity_fake = self._predict_velocity( - self.fake_model, + self.fake, renoised_xt, sigma, condition, diff --git a/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py b/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py new file mode 100644 index 000000000..06adcf1ce --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py @@ -0,0 +1,75 @@ +from lightx2v_train.model_capabilities import ( + TeacherForcingCapability, + TeacherForcingStepContext, +) +from lightx2v_train.runtime.sequence_parallel import ( + broadcast_sequence_parallel_value, +) +from lightx2v_train.schedulers.flow_matching import ( + CausalForcingFlowMatchScheduler, +) +from lightx2v_train.utils.registry import TRAINER_REGISTER + +from .base import BaseTrainer +from .flow_matching import FlowMatchingTrainer + + +@TRAINER_REGISTER("teacher_forcing") +class TeacherForcingTrainer(FlowMatchingTrainer): + trainer_name = "teacher_forcing" + required_capabilities = ( + *BaseTrainer.required_capabilities, + TeacherForcingCapability, + ) + + def __init__(self, config): + super().__init__(config) + if self.train_type != "full": + raise ValueError( + "Teacher forcing only supports training.train_type='full'." + ) + + teacher_forcing = self.training_config.get("teacher_forcing", {}) + mode = teacher_forcing.get("mode", "chunkwise") + if mode != "chunkwise": + raise ValueError( + f"Unsupported teacher_forcing.mode={mode!r}; expected " + "'chunkwise'." + ) + self.num_frame_per_chunk = int( + teacher_forcing["num_frame_per_chunk"] + ) + self.noise_augmentation_max_timestep = int( + teacher_forcing.get("noise_augmentation_max_timestep", 0) + ) + scheduler_config = self.config["scheduler"] + self.teacher_forcing_scheduler = CausalForcingFlowMatchScheduler( + num_train_timesteps=scheduler_config.get( + "num_train_timesteps", + 1000, + ), + time_shift_settings=scheduler_config.get( + "time_shift_settings", + {}, + ), + ) + + def set_model(self, model): + BaseTrainer.set_model(self, model) + self.teacher_forcing = model.capabilities.require( + TeacherForcingCapability + ) + + def compute_loss_on_sample(self, sample): + return self.teacher_forcing.compute_loss( + sample, + TeacherForcingStepContext( + scheduler=self.teacher_forcing_scheduler, + running_dtype=self.running_dtype, + num_frame_per_chunk=self.num_frame_per_chunk, + noise_augmentation_max_timestep=( + self.noise_augmentation_max_timestep + ), + broadcast=broadcast_sequence_parallel_value, + ), + ) diff --git a/lightx2v_train/lightx2v_train/trainers/tf.py b/lightx2v_train/lightx2v_train/trainers/tf.py deleted file mode 100644 index 3bfa7ebea..000000000 --- a/lightx2v_train/lightx2v_train/trainers/tf.py +++ /dev/null @@ -1,426 +0,0 @@ -import torch -import torch.distributed as dist - -from lightx2v_train.model_zoo.native.ltx2 import Modality -from lightx2v_train.runtime.distributed import get_sequence_parallel_group, get_sequence_parallel_world_size, is_sequence_parallel_enabled -from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value, sequence_parallel_frame_slice, sync_sequence_parallel_gradients -from lightx2v_train.schedulers.flow_matching import CausalForcingFlowMatchScheduler -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .flow import FlowMatchingTrainer, LTX2T2AVFlowTrainer - - -@TRAINER_REGISTER("teacher_forcing") -class TFTrainer(FlowMatchingTrainer): - trainer_name = "teacher_forcing" - - def __init__(self, config): - super().__init__(config) - if self.train_type != "full": - raise ValueError("teacher_forcing trainer only supports training.train_type='full'.") - - self.teacher_forcing_config = self.training_config.get("teacher_forcing", {}) - self.teacher_forcing_mode = self.teacher_forcing_config.get("mode", "chunkwise") - if self.teacher_forcing_mode != "chunkwise": - raise ValueError(f"Unsupported teacher_forcing.mode={self.teacher_forcing_mode!r}; only 'chunkwise' is implemented.") - - self.num_frame_per_chunk = int(self.teacher_forcing_config.get("num_frame_per_chunk")) - self.noise_augmentation_max_timestep = int(self.teacher_forcing_config.get("noise_augmentation_max_timestep", 0)) - time_shift_settings = self.config["scheduler"].get("time_shift_settings", {}) - self.teacher_forcing_scheduler = CausalForcingFlowMatchScheduler( - num_train_timesteps=self.config["scheduler"].get("num_train_timesteps", 1000), - time_shift_settings=time_shift_settings, - ) - - def compute_loss_on_sample(self, sample): - if not hasattr(self.model, "denoise_teacher_forcing"): - raise RuntimeError("teacher_forcing trainer currently requires the wan_t2v model.") - - with torch.no_grad(): - latent = self._teacher_forcing_latent(sample) - batch_size, _, num_frames, _, _ = latent.shape - sp_size = get_sequence_parallel_world_size() if is_sequence_parallel_enabled() else 1 - frame_multiple = sp_size * self.num_frame_per_chunk - if num_frames % frame_multiple != 0: - raise ValueError(f"Teacher forcing num_frames={num_frames} must be divisible by sp_size * num_frame_per_chunk = {sp_size} * {self.num_frame_per_chunk} = {frame_multiple}.") - latent = broadcast_sequence_parallel_value(latent) - - noise = torch.randn_like(latent, dtype=self.running_dtype) - noise = broadcast_sequence_parallel_value(noise) - timestep_or_sigma, timestep_weights = self.teacher_forcing_scheduler.sample_chunkwise( - batch_size=batch_size, - num_frames=num_frames, - num_frame_per_chunk=self.num_frame_per_chunk, - device=latent.device, - dtype=self.running_dtype, - ) - timestep_or_sigma = broadcast_sequence_parallel_value(timestep_or_sigma) - timestep_weights = broadcast_sequence_parallel_value(timestep_weights) - noisy_latent = self.teacher_forcing_scheduler.add_noise(latent, noise, timestep_or_sigma) - condition = self.model.encode_condition(sample) - condition = broadcast_sequence_parallel_value(condition) - - clean_latent = latent - aug_timestep_or_sigma = None - if self.noise_augmentation_max_timestep > 0: - aug_timestep_or_sigma = self.teacher_forcing_scheduler.sample_clean_augmentation( - batch_size=batch_size, - num_frames=num_frames, - num_frame_per_chunk=self.num_frame_per_chunk, - max_timestep=self.noise_augmentation_max_timestep, - device=latent.device, - dtype=self.running_dtype, - ) - aug_timestep_or_sigma = broadcast_sequence_parallel_value(aug_timestep_or_sigma) - clean_latent = self.teacher_forcing_scheduler.add_noise(latent, noise, aug_timestep_or_sigma) - - frame_start, frame_end, _ = sequence_parallel_frame_slice(num_frames, self.num_frame_per_chunk) - latent = latent[:, :, frame_start:frame_end].contiguous() - noise = noise[:, :, frame_start:frame_end].contiguous() - noisy_latent = noisy_latent[:, :, frame_start:frame_end].contiguous() - clean_latent = clean_latent[:, :, frame_start:frame_end].contiguous() - timestep_or_sigma = timestep_or_sigma[:, frame_start:frame_end].contiguous() - timestep_weights = timestep_weights[:, frame_start:frame_end].contiguous() - if aug_timestep_or_sigma is not None: - aug_timestep_or_sigma = aug_timestep_or_sigma[:, frame_start:frame_end].contiguous() - - prediction = self.model.denoise_teacher_forcing( - noisy_latent, - timestep_or_sigma, - condition, - clean_latent=clean_latent, - aug_timestep_or_sigma=aug_timestep_or_sigma, - frame_offset=frame_start, - global_num_frames=num_frames, - ) - - target = self.noise_scheduler.build_train_gt(latent, noise) - frame_loss = ((prediction.float() - target.float()) ** 2).mean(dim=(1, 3, 4)) - weighted_loss = (frame_loss * timestep_weights).sum() - if is_sequence_parallel_enabled(): - dist.all_reduce(weighted_loss, op=dist.ReduceOp.SUM, group=get_sequence_parallel_group()) - return weighted_loss / (batch_size * num_frames) - - def _teacher_forcing_latent(self, sample): - latent = sample["inputs"].get("latents") - if latent is None: - latent = self.model.encode_to_latent(sample) - latent = latent.to(device=self.model.device, dtype=self.running_dtype) - if latent.ndim == 4: - latent = latent.unsqueeze(0) - - latent_channels = self._latent_channels() - if latent_channels is not None and latent.shape[1] != latent_channels and latent.shape[2] == latent_channels: - latent = latent.permute(0, 2, 1, 3, 4).contiguous() - if latent_channels is not None and latent.shape[1] != latent_channels: - raise ValueError(f"Teacher forcing latent channels={latent.shape[1]} does not match model latent_channels={latent_channels}.") - return latent - - def _latent_channels(self): - return int(self.model_config.get("latent_channels", 16)) - - def _after_backward(self): - sync_sequence_parallel_gradients(self.trainable_params) - - -@TRAINER_REGISTER("ltx_t2av_teacher_forcing") -class LTX2T2AVTeacherForcingTrainer(LTX2T2AVFlowTrainer): - trainer_name = "ltx_t2av_teacher_forcing" - - def __init__(self, config): - super().__init__(config) - if self.train_type != "full": - raise ValueError("ltx_t2av_teacher_forcing currently supports training.train_type='full'.") - self.teacher_forcing_config = self.training_config.get("teacher_forcing", {}) - self.num_frame_per_chunk = int(self.teacher_forcing_config.get("num_frame_per_chunk", 1)) - self.noise_augmentation_max_timestep = int(self.teacher_forcing_config.get("noise_augmentation_max_timestep", 0)) - time_shift_settings = self.config["scheduler"].get("time_shift_settings", {}) - self.teacher_forcing_scheduler = CausalForcingFlowMatchScheduler( - num_train_timesteps=self.config["scheduler"].get("num_train_timesteps", 1000), - time_shift_settings=time_shift_settings, - ) - self._teacher_forcing_mask_cache = {} - - def compute_loss_on_sample(self, sample): - if is_sequence_parallel_enabled(): - raise ValueError("ltx_t2av_teacher_forcing does not support sequence parallel yet.") - - inputs = sample["inputs"] - video_data = inputs.get("video_latents") - audio_data = inputs.get("audio_latents") - conditions = sample["conditioning"].get("positive") - if video_data is None or audio_data is None or conditions is None: - raise KeyError("ltx_t2av_teacher_forcing expects inputs.video_latents, inputs.audio_latents and conditioning.positive.") - - with torch.no_grad(): - video_latents = self._video_latents(video_data) - audio_latents = self._audio_latents(audio_data) - video_latents, audio_latents = self._crop_latents_to_chunk_multiple(video_latents, audio_latents) - video_num_frames = int(video_latents.shape[2]) - video_tokens = self.video_patchifier.patchify(video_latents) - audio_tokens = self.audio_patchifier.patchify(audio_latents) - video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) - - ( - video_input, - video_target, - video_timesteps, - video_weights, - video_sigmas, - video_frame_weights, - video_aug_sigmas, - video_attention_mask, - video_block_ids, - ) = self._prepare_video_teacher_forcing_inputs( - video_tokens, - num_frames=video_num_frames, - spatial_tokens=int(video_latents.shape[3] * video_latents.shape[4]), - ) - ( - audio_input, - audio_target, - audio_timesteps, - audio_weights, - audio_sigmas, - audio_attention_mask, - audio_block_ids, - ) = self._prepare_audio_teacher_forcing_inputs( - audio_tokens, - frame_sigmas=video_sigmas, - frame_weights=video_frame_weights, - aug_frame_sigmas=video_aug_sigmas, - num_video_frames=video_num_frames, - ) - - video_positions = self._get_video_positions( - num_frames=video_num_frames, - height=int(video_data["height"][0].item()) if "height" in video_data else video_latents.shape[3], - width=int(video_data["width"][0].item()) if "width" in video_data else video_latents.shape[4], - batch_size=video_tokens.shape[0], - fps=float(video_data["fps"][0].item()) if "fps" in video_data else self.default_fps, - device=video_tokens.device, - ) - video_positions = torch.cat([video_positions, video_positions], dim=2) - audio_positions = self._get_audio_positions( - num_time_steps=audio_tokens.shape[1], - batch_size=audio_tokens.shape[0], - device=audio_tokens.device, - ) - audio_positions = torch.cat([audio_positions, audio_positions], dim=2) - video_cross_attention_mask = self._teacher_forcing_cross_attention_mask( - video_tokens.device, - video_block_ids, - audio_block_ids, - ) - audio_cross_attention_mask = self._teacher_forcing_cross_attention_mask( - audio_tokens.device, - audio_block_ids, - video_block_ids, - ) - - video_modality = Modality( - enabled=True, - latent=video_input, - sigma=video_sigmas[:, 0], - timesteps=video_timesteps, - positions=video_positions, - context=video_context, - context_mask=context_mask, - attention_mask=video_attention_mask, - cross_attention_mask=video_cross_attention_mask, - ) - audio_modality = Modality( - enabled=True, - latent=audio_input, - sigma=audio_sigmas[:, 0], - timesteps=audio_timesteps, - positions=audio_positions, - context=audio_context, - context_mask=context_mask, - attention_mask=audio_attention_mask, - cross_attention_mask=audio_cross_attention_mask, - ) - - with self.model.transformer_forward_context(): - video_pred, audio_pred = self.model.denoiser_module()( - video=video_modality, - audio=audio_modality, - perturbations=None, - ) - - video_loss = self._weighted_token_mse(video_pred, video_target, video_weights) - audio_loss = self._weighted_token_mse(audio_pred, audio_target, audio_weights) - return (self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss).mean() - - def _crop_latents_to_chunk_multiple(self, video_latents, audio_latents): - chunk = int(self.num_frame_per_chunk) - if chunk <= 1: - return video_latents, audio_latents - - num_frames = int(video_latents.shape[2]) - keep_frames = (num_frames // chunk) * chunk - if keep_frames == num_frames: - return video_latents, audio_latents - if keep_frames <= 0: - raise ValueError(f"LTX teacher forcing sample has only {num_frames} latent frames, smaller than num_frame_per_chunk={chunk}.") - - audio_frames = int(audio_latents.shape[2]) - keep_audio_frames = max(1, int(audio_frames * keep_frames // num_frames)) - video_latents = video_latents[:, :, :keep_frames].contiguous() - audio_latents = audio_latents[:, :, :keep_audio_frames].contiguous() - return video_latents, audio_latents - - def _prepare_video_teacher_forcing_inputs(self, tokens, num_frames, spatial_tokens): - batch_size = tokens.shape[0] - frame_sigmas, frame_weights = self.teacher_forcing_scheduler.sample_chunkwise( - batch_size=batch_size, - num_frames=num_frames, - num_frame_per_chunk=self.num_frame_per_chunk, - device=tokens.device, - dtype=tokens.dtype, - ) - token_sigmas = frame_sigmas.repeat_interleave(spatial_tokens, dim=1) - token_weights = frame_weights.repeat_interleave(spatial_tokens, dim=1) - noisy, target, noise = self._add_token_noise(tokens, token_sigmas) - clean = tokens - clean_timesteps = torch.zeros_like(token_sigmas) - aug_frame_sigmas = None - if self.noise_augmentation_max_timestep > 0: - aug_frame_sigmas = self.teacher_forcing_scheduler.sample_clean_augmentation( - batch_size=batch_size, - num_frames=num_frames, - num_frame_per_chunk=self.num_frame_per_chunk, - max_timestep=self.noise_augmentation_max_timestep, - device=tokens.device, - dtype=tokens.dtype, - ) - clean_timesteps = aug_frame_sigmas.repeat_interleave(spatial_tokens, dim=1) - clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise - frame_ids = torch.arange(num_frames, device=tokens.device).repeat_interleave(spatial_tokens) - block_ids = frame_ids // self.num_frame_per_chunk - attention_mask = self._teacher_forcing_attention_mask(tokens.device, block_ids) - return ( - torch.cat([clean, noisy], dim=1), - target, - torch.cat([clean_timesteps, token_sigmas], dim=1), - token_weights, - frame_sigmas, - frame_weights, - aug_frame_sigmas, - attention_mask, - block_ids, - ) - - def _prepare_audio_teacher_forcing_inputs(self, tokens, frame_sigmas, frame_weights, aug_frame_sigmas, num_video_frames): - batch_size, seq_len, _ = tokens.shape - frame_ids = torch.arange(seq_len, device=tokens.device) * max(1, int(num_video_frames)) // max(1, int(seq_len)) - frame_ids = frame_ids.clamp(max=frame_sigmas.shape[1] - 1) - token_sigmas = frame_sigmas.gather(1, frame_ids.unsqueeze(0).expand(batch_size, -1)) - token_weights = frame_weights.gather(1, frame_ids.unsqueeze(0).expand(batch_size, -1)) - noisy, target, noise = self._add_token_noise(tokens, token_sigmas) - clean = tokens - clean_timesteps = torch.zeros_like(token_sigmas) - if aug_frame_sigmas is not None: - clean_timesteps = aug_frame_sigmas.gather(1, frame_ids.unsqueeze(0).expand(batch_size, -1)) - clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise - block_ids = frame_ids // self.num_frame_per_chunk - attention_mask = self._teacher_forcing_attention_mask(tokens.device, block_ids) - return ( - torch.cat([clean, noisy], dim=1), - target, - torch.cat([clean_timesteps, token_sigmas], dim=1), - token_weights, - token_sigmas, - attention_mask, - block_ids, - ) - - def _teacher_forcing_attention_mask(self, device, block_ids): - cache_key = ("self", self._mask_device_key(device), self._block_ids_cache_key(block_ids)) - if cache_key in self._teacher_forcing_mask_cache: - return self._teacher_forcing_mask_cache[cache_key] - - transformer = self.model.denoiser_module() - for candidate in ( - transformer, - getattr(transformer, "module", None), - getattr(transformer, "_fsdp_wrapped_module", None), - ): - if candidate is None: - continue - mask_fn = getattr(candidate, "prepare_teacher_forcing_mask_from_block_ids", None) - if mask_fn is not None: - mask = mask_fn(device, block_ids) - self._teacher_forcing_mask_cache[cache_key] = mask - return mask - raise RuntimeError("ltx_t2av_teacher_forcing requires prepare_teacher_forcing_mask_from_block_ids on the LTX transformer.") - - def _teacher_forcing_cross_attention_mask(self, device, query_block_ids, key_block_ids): - cache_key = ( - "cross", - self._mask_device_key(device), - self._block_ids_cache_key(query_block_ids), - self._block_ids_cache_key(key_block_ids), - ) - if cache_key in self._teacher_forcing_mask_cache: - return self._teacher_forcing_mask_cache[cache_key] - - transformer = self.model.denoiser_module() - for candidate in ( - transformer, - getattr(transformer, "module", None), - getattr(transformer, "_fsdp_wrapped_module", None), - ): - if candidate is None: - continue - mask_fn = getattr(candidate, "prepare_teacher_forcing_cross_mask_from_block_ids", None) - if mask_fn is not None: - mask = mask_fn(device, query_block_ids, key_block_ids) - self._teacher_forcing_mask_cache[cache_key] = mask - return mask - raise RuntimeError("ltx_t2av_teacher_forcing requires prepare_teacher_forcing_cross_mask_from_block_ids on the LTX transformer.") - - @staticmethod - def _mask_device_key(device): - device = torch.device(device) - index = device.index - if device.type == "cuda" and index is None and torch.cuda.is_available(): - index = torch.cuda.current_device() - return device.type, index - - @staticmethod - def _block_ids_cache_key(block_ids): - block_ids = block_ids.detach() - if block_ids.numel() == 0: - return 0, None, None, None, None, None - block_ids_long = block_ids.to(dtype=torch.long) - positions = torch.arange(block_ids_long.numel(), device=block_ids_long.device, dtype=torch.long) - return ( - int(block_ids_long.numel()), - int(block_ids_long[0].item()), - int(block_ids_long[-1].item()), - int(block_ids_long.max().item()), - int(block_ids_long.sum().item()), - int((block_ids_long * positions).sum().item()), - ) - - @staticmethod - def _add_token_noise(tokens, sigmas, noise=None): - if noise is None: - noise = torch.randn_like(tokens) - sigmas_expanded = sigmas.unsqueeze(-1) - noisy = (1.0 - sigmas_expanded) * tokens + sigmas_expanded * noise - target = noise - tokens - return noisy, target, noise - - @staticmethod - def _causal_attention_mask(time_ids, batch_size, dtype): - allowed = time_ids.view(-1, 1) >= time_ids.view(1, -1) - return allowed.to(dtype=dtype).unsqueeze(0).expand(batch_size, -1, -1).contiguous() - - @staticmethod - def _weighted_token_mse(pred, target, weights): - pred = pred[:, -target.shape[1] :, :] - loss = (pred.float() - target.float()).pow(2).mean(dim=-1) - weights = weights.float() - return (loss * weights).sum(dim=1) / weights.sum(dim=1).clamp(min=1e-8) diff --git a/lightx2v_train/lightx2v_train/trainers/fastwam.py b/lightx2v_train/lightx2v_train/trainers/world_action.py similarity index 95% rename from lightx2v_train/lightx2v_train/trainers/fastwam.py rename to lightx2v_train/lightx2v_train/trainers/world_action.py index 2ac5a84a3..a601f7596 100644 --- a/lightx2v_train/lightx2v_train/trainers/fastwam.py +++ b/lightx2v_train/lightx2v_train/trainers/world_action.py @@ -11,6 +11,8 @@ from torch.nn.parallel import DistributedDataParallel from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, SequentialLR +from lightx2v_train.model_capabilities import WorldActionTrainingCapability +from lightx2v_train.data.utils import require_singleton_dataloader from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints from lightx2v_train.runtime.distributed import ( barrier, @@ -62,8 +64,10 @@ def _unwrap_dataset(dataset): return dataset -@TRAINER_REGISTER("fastwam") -class FastWAMTrainer: +@TRAINER_REGISTER("world_action") +class WorldActionTrainer: + required_capabilities = (WorldActionTrainingCapability,) + def __init__(self, config): self.config = config self.model_config = config["model"] @@ -119,8 +123,15 @@ def _resolve_max_train_iters(self): def set_model(self, model): self.model = model + registry = model.ensure_capabilities() + self.world_action = registry.require( + WorldActionTrainingCapability + ) def set_data(self, dataloader_train, dataloader_eval=None): + require_singleton_dataloader(dataloader_train, "Training dataloader") + if dataloader_eval is not None: + require_singleton_dataloader(dataloader_eval, "Evaluation dataloader") self.dataloader_train = dataloader_train self.dataloader_eval = dataloader_eval @@ -187,15 +198,14 @@ def _build_lr_scheduler(self): ) def setup(self, resume_ckpt_path=None): - self.model.set_dit_only_trainable() - self.model.log_model_structure() - self.trainable_params = list(self.model.trainable_parameters()) + self.world_action.configure() + self.trainable_params = list(self.world_action.parameters()) if not self.trainable_params: raise RuntimeError("FastWAM has no trainable parameters.") self.optimizer = self._build_optimizer() self.lr_scheduler = self._build_lr_scheduler() - self.train_module = self.model.unwrap_module() + self.train_module = self.world_action.module() self.ddp_enabled = False if is_distributed(): self.train_module = DistributedDataParallel( @@ -228,7 +238,7 @@ def _load_resume_state(self, resume_ckpt_path): raise RuntimeError(f"fastwam.pt not found in {resume_ckpt_path}") if not os.path.exists(state_path): raise RuntimeError(f"training_state.pt not found in {resume_ckpt_path}") - self.model.load_checkpoint(weights_path) + self.world_action.load_checkpoint(weights_path) state = torch.load(state_path, map_location="cpu", weights_only=False) if state.get("world_size") != get_world_size(): raise RuntimeError(f"Cannot resume world_size={state.get('world_size')} checkpoint with world_size={get_world_size()}.") @@ -304,7 +314,10 @@ def _save_checkpoint(self, iteration): weights_path = os.path.join(save_dir, "fastwam.pt") temporary_weights_path = _temporary_path(weights_path) try: - self.model.save_checkpoint(temporary_weights_path, step=iteration) + self.world_action.save_checkpoint( + temporary_weights_path, + step=iteration, + ) os.replace(temporary_weights_path, weights_path) finally: if os.path.exists(temporary_weights_path): @@ -334,9 +347,11 @@ def _save_checkpoint(self, iteration): ) def _forward_loss(self, sample): - with self.model.autocast_context(): - loss, loss_dict = self.train_module(sample) - return loss, loss_dict + result = self.world_action.compute_loss( + sample, + module=self.train_module, + ) + return result.loss, result.metrics def _ddp_sync_context(self, sync_grad): if self.ddp_enabled and not sync_grad: @@ -358,7 +373,7 @@ def train(self): start_time = time.perf_counter() logger.info( - "[train] start method=fastwam iter={}/{} world_size={} grad_accum={} log_every={}", + "[train] start method=world_action iter={}/{} world_size={} grad_accum={} log_every={}", current_iter, self.max_train_iters, get_world_size(), @@ -437,7 +452,7 @@ def evaluate(self, current_iter): if len(eval_dataset) == 0: return - model = self.model.unwrap_module() + model = self.world_action.module() model.eval() step_dir = os.path.join(self.eval_output_dir, f"step_{current_iter:09d}") os.makedirs(step_dir, exist_ok=True) @@ -492,7 +507,7 @@ def evaluate(self, current_iter): ) barrier() - self.model.set_dit_only_trainable() + self.world_action.configure() def _to_batched_eval_sample(self, sample): result = {} @@ -521,8 +536,9 @@ def _evaluate_sample(self, model, dataset, sample, sample_number, eval_index, st rng_devices = [torch.cuda.current_device()] if torch.cuda.is_available() else [] with torch.random.fork_rng(devices=rng_devices): torch.manual_seed(sample_seed) - with self.model.autocast_context(): - val_loss, val_loss_dict = model.training_loss(sample) + evaluation = self.world_action.evaluation_loss(sample) + val_loss = evaluation.loss + val_loss_dict = evaluation.metrics result = { "sample_number": int(sample_number), "dataset_index": int(eval_index), diff --git a/lightx2v_train/lightx2v_train/tricks/cdm.py b/lightx2v_train/lightx2v_train/tricks/cdm.py index 648f6172f..672495523 100644 --- a/lightx2v_train/lightx2v_train/tricks/cdm.py +++ b/lightx2v_train/lightx2v_train/tricks/cdm.py @@ -14,7 +14,6 @@ class CdmScheduler(Protocol): def sigma_at( self, step_index: int, - batch_size: int, *, device: torch.device, dtype: torch.dtype, @@ -22,7 +21,6 @@ def sigma_at( def sample_renoise_sigma( self, - batch_size: int, *, device: torch.device, dtype: torch.dtype, @@ -71,8 +69,6 @@ def from_mapping(cls, config: Mapping[str, Any]) -> "CdmConfig": @dataclass(frozen=True) class CdmTrainerConstraints: supported: bool - model_name: str - supported_model_names: frozenset[str] @dataclass(frozen=True) @@ -133,9 +129,6 @@ def validate_trainer( return if not constraints.supported: raise ValueError(f"{trainer_name} does not support training.dmd.cdm.") - if constraints.model_name not in constraints.supported_model_names: - supported = ", ".join(repr(name) for name in sorted(constraints.supported_model_names)) - raise ValueError(f"training.dmd.cdm is only supported for model.name in {{{supported}}}, got {constraints.model_name!r}.") def effective_weight(self, current_iteration: int | None) -> float: if self.config.warmup_iters <= 0 or current_iteration is None: @@ -167,15 +160,12 @@ def student_loss( }, ) - batch_size = context.trajectory_latent.shape[0] trajectory_sigma = context.scheduler.sigma_at( context.end_step_index, - batch_size, device=context.device, dtype=torch.float32, ) student_sigma = context.scheduler.sample_renoise_sigma( - batch_size, device=context.device, dtype=torch.float32, ) @@ -196,7 +186,6 @@ def student_loss( student_x0 = (student_xt - student_sigma_expanded * student_prediction).to(context.dtype) teacher_sigma = context.scheduler.sample_renoise_sigma( - batch_size, device=context.device, dtype=context.dtype, ) diff --git a/lightx2v_train/lightx2v_train/tricks/diversity.py b/lightx2v_train/lightx2v_train/tricks/diversity.py index 9cd7604b8..dd5e2a4e6 100644 --- a/lightx2v_train/lightx2v_train/tricks/diversity.py +++ b/lightx2v_train/lightx2v_train/tricks/diversity.py @@ -23,7 +23,6 @@ def set_timesteps( def sigma_at( self, step_index: int, - batch_size: int, *, device: torch.device, dtype: torch.dtype, @@ -180,7 +179,6 @@ def _compute_anchor( for step_index in range(self.config.anchor_step): sigma = scheduler.sigma_at( step_index, - context.initial_noise.shape[0], device=context.device, dtype=context.dtype, ) @@ -197,7 +195,6 @@ def _compute_anchor( ) anchor_sigma = scheduler.sigma_at( self.config.anchor_step, - context.initial_noise.shape[0], device=context.device, dtype=context.dtype, ) @@ -218,7 +215,6 @@ def student_loss( anchor_latent, anchor_sigma = self._compute_anchor(context) sigma_first = context.student_scheduler.sigma_at( 0, - initial_noise.shape[0], device=context.device, dtype=context.dtype, ) diff --git a/lightx2v_train/lightx2v_train/tricks/real_data_fake.py b/lightx2v_train/lightx2v_train/tricks/real_data_fake.py index ebe96600d..ef26b7352 100644 --- a/lightx2v_train/lightx2v_train/tricks/real_data_fake.py +++ b/lightx2v_train/lightx2v_train/tricks/real_data_fake.py @@ -241,12 +241,11 @@ def _euler_step( def _phase_sigma( self, - batch_size: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: raw_timestep = torch.full( - (int(batch_size),), + (1,), self.setup_context.match_timestep, device=device, dtype=torch.long, @@ -270,7 +269,6 @@ def _build_student_outputs( ) if setup.mode == "phased" and region == "high": sigma_s = self._phase_sigma( - xt.shape[0], xt.device, xt.dtype, ) @@ -292,7 +290,6 @@ def _build_student_outputs( def _sample_score_sigma( self, context: RealDataFakeStepContext, - batch_size: int, ) -> torch.Tensor: setup = self.setup_context if setup.mode == "standard": @@ -306,7 +303,7 @@ def _sample_score_sigma( timestep = torch.randint( min_timestep, max_timestep, - (int(batch_size),), + (1,), device=context.device, dtype=torch.long, ).float() @@ -343,7 +340,7 @@ def _sample_score_sigma( candidate_indices = torch.randint( 0, candidate_sigmas.numel(), - (int(batch_size),), + (1,), device=context.device, dtype=torch.long, ) @@ -423,12 +420,14 @@ def _real_state( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: config = self._region_config(context.region) real_latents = context.extract_real_latents(context.sample) + if real_latents.shape[0] != 1: + raise ValueError("Real-data DMD only supports physical batch size 1.") timestep_index = context.sample_synced_int( 0, len(config.timestep_list), ) raw_timestep = torch.full( - (real_latents.shape[0],), + (1,), config.timestep_list[timestep_index], device=context.device, dtype=torch.long, @@ -469,7 +468,7 @@ def student_loss( context, grad_enabled=True, ) - sigma_t = self._sample_score_sigma(context, anchor.shape[0]) + sigma_t = self._sample_score_sigma(context) noise = context.broadcast_noise(torch.randn_like(anchor, dtype=torch.float32)) with torch.no_grad(): score_xt = self._score_forward( @@ -523,7 +522,7 @@ def fake_loss( context, grad_enabled=False, ) - sigma_t = self._sample_score_sigma(context, anchor.shape[0]) + sigma_t = self._sample_score_sigma(context) noise = context.broadcast_noise(torch.randn_like(anchor, dtype=torch.float32)) with torch.no_grad(): score_xt = self._score_forward( diff --git a/lightx2v_train/lightx2v_train/utils/image_size_buckets.py b/lightx2v_train/lightx2v_train/utils/image_size_buckets.py new file mode 100644 index 000000000..be9a33a5b --- /dev/null +++ b/lightx2v_train/lightx2v_train/utils/image_size_buckets.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from numbers import Integral, Real + + +@dataclass(frozen=True) +class ImageSizeBucket: + value: tuple[int, ...] + ratio: float | None + + @property + def spatial_size(self) -> tuple[int, int]: + return self.value[-2], self.value[-1] + + +def parse_image_size_buckets( + entries, + *, + config_path: str = "training.dmd.image_sizes", +) -> list[ImageSizeBucket]: + """Parse the only supported image-size bucket schemas. + + Every entry is either ``{"value": [...]}`` or + ``{"value": [...], "ratio": number}``. A configuration cannot mix the + weighted and unweighted forms. The final two values are always interpreted + as pixel height and width. + """ + if entries is None: + return [] + if not isinstance(entries, list): + raise TypeError(f"{config_path} must be a list, got {type(entries).__name__}.") + + buckets = [] + ratio_modes = set() + spatial_sizes = set() + for index, entry in enumerate(entries): + entry_path = f"{config_path}[{index}]" + if not isinstance(entry, Mapping): + raise TypeError( + f"{entry_path} must be {{'value': [...]}} or " + "{'value': [...], 'ratio': number}; legacy bare lists are not supported." + ) + keys = set(entry) + if keys not in ({"value"}, {"value", "ratio"}): + raise ValueError( + f"{entry_path} must contain exactly 'value' or 'value' and 'ratio', " + f"got keys {sorted(keys)}." + ) + + value = entry["value"] + if not isinstance(value, list) or len(value) not in {2, 3}: + raise ValueError( + f"{entry_path}.value must be [height, width] or " + f"[prefix, height, width], got {value!r}." + ) + if any(not isinstance(item, Integral) or isinstance(item, bool) for item in value): + raise TypeError(f"{entry_path}.value must contain integers, got {value!r}.") + normalized_value = tuple(int(item) for item in value) + if any(item <= 0 for item in normalized_value): + raise ValueError(f"{entry_path}.value must contain positive integers, got {value!r}.") + + has_ratio = "ratio" in entry + ratio_modes.add(has_ratio) + ratio = None + if has_ratio: + raw_ratio = entry["ratio"] + if not isinstance(raw_ratio, Real) or isinstance(raw_ratio, bool): + raise TypeError(f"{entry_path}.ratio must be a positive number, got {raw_ratio!r}.") + ratio = float(raw_ratio) + if not math.isfinite(ratio) or ratio <= 0: + raise ValueError(f"{entry_path}.ratio must be finite and positive, got {raw_ratio!r}.") + + bucket = ImageSizeBucket(value=normalized_value, ratio=ratio) + if bucket.spatial_size in spatial_sizes: + height, width = bucket.spatial_size + raise ValueError( + f"{config_path} contains duplicate spatial bucket {height}x{width}." + ) + spatial_sizes.add(bucket.spatial_size) + buckets.append(bucket) + + if len(ratio_modes) > 1: + raise ValueError( + f"{config_path} cannot mix entries with and without ratio; " + "use one schema consistently." + ) + return buckets diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index d202c454f..1ae449caa 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -72,40 +72,33 @@ def merge(self, other_register): _MODEL_MODULES = { - "flux2_dev": "lightx2v_train.model_zoo.flux2_dev", - "flux2_klein": "lightx2v_train.model_zoo.flux2_klein", - "lingbot_video": "lightx2v_train.model_zoo.lingbot_video", - "longcat_image": "lightx2v_train.model_zoo.longcat_image", - "ltx_t2av": "lightx2v_train.model_zoo.ltx_t2av", - "ltx_t2av_ar": "lightx2v_train.model_zoo.ltx_t2av", - "qwen_image": "lightx2v_train.model_zoo.qwen_image", - "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image_edit", - "wan_t2v": "lightx2v_train.model_zoo.wan_t2v", - "wan_t2v_ar": "lightx2v_train.model_zoo.wan_t2v", - "wan_t2v_14b": "lightx2v_train.model_zoo.wan_t2v", - "wan_t2v_14b_ar": "lightx2v_train.model_zoo.wan_t2v", - "wan_fastwam": "lightx2v_train.model_zoo.wan_fastwam", - "wan_ti2v_5b": "lightx2v_train.model_zoo.wan_ti2v_5b", - "wan_ti2v_5b_ar": "lightx2v_train.model_zoo.wan_ti2v_5b", + "flux2_dev": "lightx2v_train.model_zoo.flux2.flux2_dev", + "flux2_klein": "lightx2v_train.model_zoo.flux2.flux2_klein", + "lingbot_video": "lightx2v_train.model_zoo.wan.lingbot_video", + "longcat_image": "lightx2v_train.model_zoo.longcat_image.longcat_image", + "ltx_t2av": "lightx2v_train.model_zoo.ltx.ltx_t2av", + "ltx_t2av_ar": "lightx2v_train.model_zoo.ltx.ltx_t2av", + "qwen_image": "lightx2v_train.model_zoo.qwen_image.qwen_image", + "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image.qwen_image_edit", + "wan_t2v": "lightx2v_train.model_zoo.wan.wan_t2v", + "wan_t2v_ar": "lightx2v_train.model_zoo.wan.wan_t2v", + "wan_t2v_14b": "lightx2v_train.model_zoo.wan.wan_t2v", + "wan_t2v_14b_ar": "lightx2v_train.model_zoo.wan.wan_t2v", + "wan_fastwam": "lightx2v_train.model_zoo.wan.wan_fastwam", + "wan_ti2v_5b": "lightx2v_train.model_zoo.wan.wan_ti2v_5b", + "wan_ti2v_5b_ar": "lightx2v_train.model_zoo.wan.wan_ti2v_5b", } _TRAINER_MODULES = { - "cm": "lightx2v_train.trainers.consistency.trainer", + "autoregressive_dmd": "lightx2v_train.trainers.dmd.autoregressive_dmd", "consistency": "lightx2v_train.trainers.consistency.trainer", "dmd": "lightx2v_train.trainers.dmd.trainer", "dopsd": "lightx2v_train.trainers.dopsd", - "fastwam": "lightx2v_train.trainers.fastwam", - "flow": "lightx2v_train.trainers.flow", - "lingbot_video_dmd": "lightx2v_train.trainers.dmd.video_trainer", - "ltx_t2av_ar_dmd": "lightx2v_train.trainers.dmd.ltx_trainer", - "ltx_t2av_dmd": "lightx2v_train.trainers.dmd.ltx_trainer", - "ltx_t2av_flow": "lightx2v_train.trainers.flow", - "ltx_t2av_teacher_forcing": "lightx2v_train.trainers.tf", - "teacher_forcing": "lightx2v_train.trainers.tf", - "video_ar_dmd": "lightx2v_train.trainers.dmd.video_ar_trainer", - "video_dmd": "lightx2v_train.trainers.dmd.video_trainer", - "video_phased_dmd": "lightx2v_train.trainers.phased_dmd.trainer", - "video_sgmd": "lightx2v_train.trainers.sgmd", + "flow_matching": "lightx2v_train.trainers.flow_matching", + "phased_dmd": "lightx2v_train.trainers.phased_dmd.trainer", + "sgmd": "lightx2v_train.trainers.sgmd", + "teacher_forcing": "lightx2v_train.trainers.teacher_forcing", + "world_action": "lightx2v_train.trainers.world_action", } _INFERENCER_MODULES = { @@ -173,8 +166,12 @@ def build_data(config, train_or_val): if train_or_val not in data_config: available_splits = ", ".join(repr(k) for k in sorted(data_config.keys())) raise ValueError(f"config['data'] has no key {train_or_val!r}. Available keys: {available_splits}") - data_config_split = data_config[train_or_val] + data_config_split = dict(data_config[train_or_val]) data_name = data_config_split.get("name", "image_dataset") + if train_or_val == "train" and data_name == "prompt_dataset": + image_sizes = config.get("training", {}).get("dmd", {}).get("image_sizes") + if image_sizes is not None: + data_config_split["image_sizes"] = image_sizes _ensure_data_registered(data_name) if data_name not in DATA_REGISTER: available_names = ", ".join(sorted(DATA_REGISTER.keys())) From 23b5925b84fd3d37e5f38f398abe437664ddc4c0 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 10 Aug 2026 07:58:04 +0000 Subject: [PATCH 04/34] fix ci --- .../lightx2v_train/model_capabilities/base.py | 19 ++------ .../model_zoo/capability_adapters/common.py | 46 ++++--------------- .../ltx_distillation_capability.py | 4 +- .../wan_distillation_capability.py | 5 +- .../trainers/teacher_forcing.py | 25 +++------- .../lightx2v_train/trainers/world_action.py | 6 +-- .../utils/image_size_buckets.py | 24 ++-------- 7 files changed, 28 insertions(+), 101 deletions(-) diff --git a/lightx2v_train/lightx2v_train/model_capabilities/base.py b/lightx2v_train/lightx2v_train/model_capabilities/base.py index 565eb2eb3..b13dc8795 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/base.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/base.py @@ -26,14 +26,9 @@ def register( implementation: CapabilityT, ) -> None: if not isinstance(implementation, interface): - raise TypeError( - f"{type(implementation).__name__} does not implement " - f"{interface.__name__}." - ) + raise TypeError(f"{type(implementation).__name__} does not implement {interface.__name__}.") if interface in self._capabilities: - raise RuntimeError( - f"Capability already registered: {interface.__name__}." - ) + raise RuntimeError(f"Capability already registered: {interface.__name__}.") self._capabilities[interface] = implementation def supports(self, interface: type[ModelCapability]) -> bool: @@ -42,20 +37,14 @@ def supports(self, interface: type[ModelCapability]) -> bool: def require(self, interface: type[CapabilityT]) -> CapabilityT: implementation = self._capabilities.get(interface) if implementation is None: - raise CapabilityNotSupportedError( - f"Missing capability: {interface.__name__}." - ) + 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) - ) + return tuple(interface for interface in interfaces if not self.supports(interface)) class CapabilityProvider: diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py index 99a2a2b9a..235088366 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -30,10 +30,7 @@ def _require_single_prompt(prompt): return True prompts = list(prompt) if len(prompts) != 1: - raise ValueError( - "Training requires exactly one prompt per rank; " - f"physical batch sizes greater than 1 are not supported, got {len(prompts)} prompts." - ) + 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 @@ -50,10 +47,7 @@ def _negative_prompt(conditioning, fallback, scalar=False): elif len(prompts) == 1: prompt = prompts[0] else: - raise ValueError( - "Training requires exactly one negative prompt per rank; " - f"got {len(prompts)} negative prompts." - ) + 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] @@ -61,10 +55,7 @@ def _negative_prompt(conditioning, fallback, scalar=False): 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 " - f"than 1 are not supported, got {shape}." - ) + raise ValueError(f"{name} must have leading dimension 1; physical batch sizes greater than 1 are not supported, got {shape}.") return value @@ -269,38 +260,21 @@ def latent_shape( width = int(broadcast(width)) if image_sizes: - configured_buckets = { - bucket.spatial_size - for bucket in parse_image_size_buckets(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( - "Image DMD sample size " - f"{height}x{width} is not in training.dmd.image_sizes: " - f"[{configured}]." - ) + 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}].") spatial_scale = int(self.model.vae_scale_factor) size_multiple = spatial_scale * 2 if height % size_multiple or width % size_multiple: - raise ValueError( - "Image DMD target_height and target_width must be divisible by " - f"the model image size multiple {size_multiple}, got " - f"{height}x{width}." - ) + raise ValueError(f"Image DMD target_height and target_width must be divisible by the model image size multiple {size_multiple}, got {height}x{width}.") return self.model.dmd_latent_shape(height, width) @staticmethod def _target_dimension(meta, key): if key not in meta: - raise KeyError( - "Image DMD prompt-only samples require meta.target_height and " - f"meta.target_width; missing {key}. target_image is not used." - ) + 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() @@ -310,9 +284,7 @@ def _target_dimension(meta, key): 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}." - ) + 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] diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py index f05f3e8f1..1b5810306 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py @@ -433,8 +433,6 @@ def _singleton_shape(shape, name): return (1, *shape) if len(shape) == expected: if shape[0] != 1: - raise ValueError( - f"training.dmd.{name} must begin with singleton dimension 1, got {shape}." - ) + raise ValueError(f"training.dmd.{name} must begin with singleton dimension 1, got {shape}.") return shape raise ValueError(f"training.dmd.{name} has unsupported shape {shape}.") diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py index 466628602..faf3623f4 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py @@ -39,10 +39,7 @@ def latent_shape( if configured is not None: shape = tuple(int(dimension) for dimension in configured) if len(shape) != 5 or shape[0] != 1: - raise ValueError( - "training.dmd.image_or_video_shape must be a five-dimensional " - f"singleton shape beginning with 1, got {shape}." - ) + raise ValueError(f"training.dmd.image_or_video_shape must be a five-dimensional singleton shape beginning with 1, got {shape}.") return shape missing = [key for key in ("height", "width", "num_frames") if key not in shape_config] if missing: diff --git a/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py b/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py index 06adcf1ce..252b97f6c 100644 --- a/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py +++ b/lightx2v_train/lightx2v_train/trainers/teacher_forcing.py @@ -25,23 +25,14 @@ class TeacherForcingTrainer(FlowMatchingTrainer): def __init__(self, config): super().__init__(config) if self.train_type != "full": - raise ValueError( - "Teacher forcing only supports training.train_type='full'." - ) + raise ValueError("Teacher forcing only supports training.train_type='full'.") teacher_forcing = self.training_config.get("teacher_forcing", {}) mode = teacher_forcing.get("mode", "chunkwise") if mode != "chunkwise": - raise ValueError( - f"Unsupported teacher_forcing.mode={mode!r}; expected " - "'chunkwise'." - ) - self.num_frame_per_chunk = int( - teacher_forcing["num_frame_per_chunk"] - ) - self.noise_augmentation_max_timestep = int( - teacher_forcing.get("noise_augmentation_max_timestep", 0) - ) + raise ValueError(f"Unsupported teacher_forcing.mode={mode!r}; expected 'chunkwise'.") + self.num_frame_per_chunk = int(teacher_forcing["num_frame_per_chunk"]) + self.noise_augmentation_max_timestep = int(teacher_forcing.get("noise_augmentation_max_timestep", 0)) scheduler_config = self.config["scheduler"] self.teacher_forcing_scheduler = CausalForcingFlowMatchScheduler( num_train_timesteps=scheduler_config.get( @@ -56,9 +47,7 @@ def __init__(self, config): def set_model(self, model): BaseTrainer.set_model(self, model) - self.teacher_forcing = model.capabilities.require( - TeacherForcingCapability - ) + self.teacher_forcing = model.capabilities.require(TeacherForcingCapability) def compute_loss_on_sample(self, sample): return self.teacher_forcing.compute_loss( @@ -67,9 +56,7 @@ def compute_loss_on_sample(self, sample): scheduler=self.teacher_forcing_scheduler, running_dtype=self.running_dtype, num_frame_per_chunk=self.num_frame_per_chunk, - noise_augmentation_max_timestep=( - self.noise_augmentation_max_timestep - ), + noise_augmentation_max_timestep=(self.noise_augmentation_max_timestep), broadcast=broadcast_sequence_parallel_value, ), ) diff --git a/lightx2v_train/lightx2v_train/trainers/world_action.py b/lightx2v_train/lightx2v_train/trainers/world_action.py index a601f7596..819e8ca01 100644 --- a/lightx2v_train/lightx2v_train/trainers/world_action.py +++ b/lightx2v_train/lightx2v_train/trainers/world_action.py @@ -11,8 +11,8 @@ from torch.nn.parallel import DistributedDataParallel from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, SequentialLR -from lightx2v_train.model_capabilities import WorldActionTrainingCapability from lightx2v_train.data.utils import require_singleton_dataloader +from lightx2v_train.model_capabilities import WorldActionTrainingCapability from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints from lightx2v_train.runtime.distributed import ( barrier, @@ -124,9 +124,7 @@ def _resolve_max_train_iters(self): def set_model(self, model): self.model = model registry = model.ensure_capabilities() - self.world_action = registry.require( - WorldActionTrainingCapability - ) + self.world_action = registry.require(WorldActionTrainingCapability) def set_data(self, dataloader_train, dataloader_eval=None): require_singleton_dataloader(dataloader_train, "Training dataloader") diff --git a/lightx2v_train/lightx2v_train/utils/image_size_buckets.py b/lightx2v_train/lightx2v_train/utils/image_size_buckets.py index be9a33a5b..f3dfbda5c 100644 --- a/lightx2v_train/lightx2v_train/utils/image_size_buckets.py +++ b/lightx2v_train/lightx2v_train/utils/image_size_buckets.py @@ -39,23 +39,14 @@ def parse_image_size_buckets( for index, entry in enumerate(entries): entry_path = f"{config_path}[{index}]" if not isinstance(entry, Mapping): - raise TypeError( - f"{entry_path} must be {{'value': [...]}} or " - "{'value': [...], 'ratio': number}; legacy bare lists are not supported." - ) + raise TypeError(f"{entry_path} must be {{'value': [...]}} or {{'value': [...], 'ratio': number}}; legacy bare lists are not supported.") keys = set(entry) if keys not in ({"value"}, {"value", "ratio"}): - raise ValueError( - f"{entry_path} must contain exactly 'value' or 'value' and 'ratio', " - f"got keys {sorted(keys)}." - ) + raise ValueError(f"{entry_path} must contain exactly 'value' or 'value' and 'ratio', got keys {sorted(keys)}.") value = entry["value"] if not isinstance(value, list) or len(value) not in {2, 3}: - raise ValueError( - f"{entry_path}.value must be [height, width] or " - f"[prefix, height, width], got {value!r}." - ) + raise ValueError(f"{entry_path}.value must be [height, width] or [prefix, height, width], got {value!r}.") if any(not isinstance(item, Integral) or isinstance(item, bool) for item in value): raise TypeError(f"{entry_path}.value must contain integers, got {value!r}.") normalized_value = tuple(int(item) for item in value) @@ -76,15 +67,10 @@ def parse_image_size_buckets( bucket = ImageSizeBucket(value=normalized_value, ratio=ratio) if bucket.spatial_size in spatial_sizes: height, width = bucket.spatial_size - raise ValueError( - f"{config_path} contains duplicate spatial bucket {height}x{width}." - ) + raise ValueError(f"{config_path} contains duplicate spatial bucket {height}x{width}.") spatial_sizes.add(bucket.spatial_size) buckets.append(bucket) if len(ratio_modes) > 1: - raise ValueError( - f"{config_path} cannot mix entries with and without ratio; " - "use one schema consistently." - ) + raise ValueError(f"{config_path} cannot mix entries with and without ratio; use one schema consistently.") return buckets From 5aa79454689f77f5cbbdceae864ccaf1d9b829f8 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 11 Aug 2026 10:58:14 +0000 Subject: [PATCH 05/34] update --- .../train/flow/wan2_1_t2v_1_3b_lora.yaml | 13 +++++------ lightx2v_train/scripts/run_wan_t2v.sh | 22 +++++-------------- 2 files changed, 12 insertions(+), 23 deletions(-) mode change 100644 => 100755 lightx2v_train/scripts/run_wan_t2v.sh 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 a7962a71d..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 @@ -50,8 +49,8 @@ data: val: name: prompt_dataset data_path: - - /path/to/wan_t2v_prompts.txt - max_samples: 5 + - /root/yongyang/datasets/LightX2V_train_data_examples/dataset_v2/val.jsonl + max_samples: 4 num_workers: 0 shuffle: false diff --git a/lightx2v_train/scripts/run_wan_t2v.sh b/lightx2v_train/scripts/run_wan_t2v.sh old mode 100644 new mode 100755 index 36a856709..5224d5e6b --- a/lightx2v_train/scripts/run_wan_t2v.sh +++ b/lightx2v_train/scripts/run_wan_t2v.sh @@ -1,18 +1,8 @@ -#!/bin/bash +#!/usr/bin/env bash -set -euo pipefail +export CUDA_VISIBLE_DEVICES=0 +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} -lightx2v_path=${LIGHTX2V_PATH:-/path/to/LightX2V} -diffusers_path=${DIFFUSERS_PATH:-/path/to/diffusers/src} -config=${CONFIG:-${lightx2v_path}/lightx2v_train/configs/train/tf/wan2_1_t2v_1_3b_tf_chunkwise_full.yaml} - -cd "${lightx2v_path}/lightx2v_train" -export PYTHONPATH="${diffusers_path}:${lightx2v_path}:${PYTHONPATH:-}" -export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-3} - -nproc_per_node=${NPROC_PER_NODE:-1} - -torchrun \ ---standalone \ ---nproc_per_node="${nproc_per_node}" \ -train.py --config "${config}" +torchrun --standalone --nproc_per_node=1 \ + train.py \ + --config configs/train/flow/wan2_1_t2v_1_3b_lora.yaml From eb250fae7b5d45e2fcf9fc8e5c3296d2e91b61da Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 11 Aug 2026 12:08:39 +0000 Subject: [PATCH 06/34] 1 --- lightx2v_train/scripts/run_wan_t2v.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightx2v_train/scripts/run_wan_t2v.sh b/lightx2v_train/scripts/run_wan_t2v.sh index 5224d5e6b..9b1870f38 100755 --- a/lightx2v_train/scripts/run_wan_t2v.sh +++ b/lightx2v_train/scripts/run_wan_t2v.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash export CUDA_VISIBLE_DEVICES=0 -export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --standalone --nproc_per_node=1 \ train.py \ From e3e2445c06e357fb8f7b2ac7553d52fd3c158d75 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 17 Aug 2026 12:14:24 +0000 Subject: [PATCH 07/34] update minimax h3 --- .../train/dmd/minimax_h3_t2av_dmd_lora.yaml | 6 + .../lightx2v_train/model_zoo/__init__.py | 2 +- .../model_zoo/minimax_h3/__init__.py | 5 + .../capability_adapters/__init__.py | 10 + .../minimax_h3/capability_adapters/common.py | 27 ++ .../minimax_h3_distillation_capability.py | 339 ++++++++++++++++++ .../{ => minimax_h3}/minimax_h3_t2av.py | 50 ++- .../trainers/dmd/minimax_h3_trainer.py | 279 +++----------- .../lightx2v_train/utils/registry.py | 2 +- .../scripts/run_minimax_h3_t2av_dmd.sh | 14 + 10 files changed, 498 insertions(+), 236 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/minimax_h3/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/common.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py rename lightx2v_train/lightx2v_train/model_zoo/{ => minimax_h3}/minimax_h3_t2av.py (66%) create mode 100755 lightx2v_train/scripts/run_minimax_h3_t2av_dmd.sh 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..170ec9171 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,6 +10,8 @@ model: video_latent_channels: 24 audio_latent_channels: 32 vae_spatial_scale_factor: 16 + # 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 @@ -57,6 +61,7 @@ training: save_every_iters: 100 save_total_limit: 5 minimax_h3: + # These shifts are applied independently inside the H3 capability. video_flow_shift: 6.0 audio_flow_shift: 3.0 video_loss_weight: 1.0 @@ -72,6 +77,7 @@ training: num_frames: 124 renoise_sigma_min: 0.02 renoise_sigma_max: 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/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index ad0890562..cec809811 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -8,7 +8,7 @@ "LingBotVideoModel": (".wan.lingbot_video", "LingBotVideoModel"), "LongCatImageModel": (".longcat_image.longcat_image", "LongCatImageModel"), "LTX2T2AVModel": (".ltx.ltx_t2av", "LTX2T2AVModel"), - "MiniMaxH3T2AVModel": (".minimax_h3_t2av", "MiniMaxH3T2AVModel"), + "MiniMaxH3T2AVModel": (".minimax_h3.minimax_h3_t2av", "MiniMaxH3T2AVModel"), "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), "QwenImageEditModel": (".qwen_image.qwen_image_edit", "QwenImageEditModel"), "WanFastWAMModel": (".wan.wan_fastwam", "WanFastWAMModel"), 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..2788ef6e3 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/__init__.py @@ -0,0 +1,10 @@ +"""Capability adapters for MiniMax-H3.""" + +from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape +from .minimax_h3_distillation_capability import MiniMaxH3DistillationCapability + +__all__ = [ + "MiniMaxH3DistillationCapability", + "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_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py new file mode 100644 index 000000000..0f6d63d1b --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py @@ -0,0 +1,339 @@ +"""Distribution-matching capability for MiniMax-H3 T2AV.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistillationCapability, + _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 + + +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 MiniMaxH3DistillationCapability(GenericDistillationCapability): + """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", + ) + + def __init__(self, model) -> None: + super().__init__(model) + config = model.config["training"].get("minimax_h3", {}) + self.video_weight = float(config.get("video_loss_weight", 1.0)) + self.audio_weight = float(config.get("audio_loss_weight", 1.0)) + self.video_shift = float(config.get("video_flow_shift", 12.0)) + self.audio_shift = float(config.get("audio_flow_shift", 3.0)) + if self.video_shift <= 0 or self.audio_shift <= 0: + raise ValueError("MiniMax-H3 video_flow_shift and audio_flow_shift must be positive.") + if self.video_weight < 0 or self.audio_weight < 0: + raise ValueError("MiniMax-H3 video_loss_weight and audio_loss_weight cannot be negative.") + if self.video_weight == 0 and self.audio_weight == 0: + raise ValueError("At least one MiniMax-H3 modality loss weight must be non-zero.") + self._layout_cache = {} + + @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.") + positive = batch["conditioning"].get("positive") + if positive is None: + raise KeyError("MiniMax-H3 DMD requires conditioning.positive with cached prompt_embeds and text_token_tags. Build it with data_process/minimax_h3/build_minimax_h3_prompt_latents.py.") + with torch.no_grad(): + condition = self.model.prepare_text_condition(positive) + 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_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py similarity index 66% rename from lightx2v_train/lightx2v_train/model_zoo/minimax_h3_t2av.py rename to lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py index f4eb6ca3c..0b3d4d066 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3_t2av.py +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/minimax_h3_t2av.py @@ -5,28 +5,41 @@ import torch from peft import LoraConfig, inject_adapter_in_model +from lightx2v_train.model_capabilities import DistillationCapability +from lightx2v_train.model_zoo.minimax_h3.capability_adapters import MiniMaxH3DistillationCapability +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 .native.minimax_h3 import load_minimax_h3_transformer +from ..base import BaseModel @MODEL_REGISTER("minimax_h3_t2av") class MiniMaxH3T2AVModel(BaseModel): - """A normal ``nn.Module`` H3 model suitable for LoRA/full DMD training.""" + """A standard trainable wrapper around Diffusers' MiniMax-H3 module.""" pipeline_cls = None + def register_capabilities(self): + super().register_capabilities() + self.capabilities.register( + DistillationCapability, + MiniMaxH3DistillationCapability(self), + ) + 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.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)) + if self.video_latent_channels <= 0 or self.audio_latent_channels <= 0 or self.vae_spatial_scale_factor <= 0: + raise ValueError("MiniMax-H3 latent channels and VAE spatial scale must be positive.") self.use_autocast = bool(config.get("use_autocast", False)) self.transformer = load_minimax_h3_transformer( self.pretrained_model_path, @@ -46,7 +59,7 @@ def transformer_forward_context(self): 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"] + target_modules = MiniMaxH3DistillationCapability._DEFAULT_LORA_TARGETS lora_config = LoraConfig( r=rank, lora_alpha=alpha, @@ -54,17 +67,26 @@ def add_lora(self, rank, alpha, target_modules): target_modules=target_modules, ) try: - self.transformer = inject_adapter_in_model(lora_config, self.transformer, adapter_name="default") + 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.") + 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: @@ -82,7 +104,7 @@ def prepare_text_condition(self, condition): 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.") + raise RuntimeError("MiniMax-H3 DMD expects cached prompt latents. Run data_process/minimax_h3/build_minimax_h3_prompt_latents.py first.") def enable_gradient_checkpointing(self): if hasattr(self.transformer, "enable_gradient_checkpointing"): @@ -94,6 +116,12 @@ 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)}, + { + "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/trainers/dmd/minimax_h3_trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py index ec599c00b..7f5cd9d1e 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py @@ -1,246 +1,79 @@ -"""DMD distillation for MiniMax-H3 text-to-audio-video generation.""" +"""Thin DMD trainer specialization for MiniMax-H3 T2AV.""" import torch from loguru import logger -from lightx2v_train.model_zoo.native.minimax_h3 import ( - audio_latent_num_frames, - build_packed_sequence, - build_row_timesteps, - video_latent_num_frames, -) from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value from lightx2v_train.utils.registry import TRAINER_REGISTER -from .math import dmd_loss_pair, weighted_mse_pair -from .video_trainer import VideoDmdTrainer - - -def _shift_sigma(sigma: torch.Tensor, shift: float) -> torch.Tensor: - return shift * sigma / (1.0 + (shift - 1.0) * sigma) +from .trainer import DmdTrainer @TRAINER_REGISTER("minimax_h3_t2av_dmd") -class MiniMaxH3T2AVDmdTrainer(VideoDmdTrainer): - """Joint video/audio DMD with H3's data-ward velocity convention.""" +class MiniMaxH3T2AVDmdTrainer(DmdTrainer): + """Configure generic DMD for H3's supported training feature set. + + Audio/video tensor math and the clean-ward velocity convention live in + ``MiniMaxH3DistillationCapability``. This class only validates algorithm + options and supplies H3's continuous, unshifted score-sigma sampling. + """ trainer_name = "minimax_h3_t2av_dmd" - allowed_model_names = {"minimax_h3_t2av"} default_negative_prompt = "" - default_lora_target_modules = ("to_q", "to_k", "to_v", "to_out.0", "ff.net.0.proj", "ff.net.2") + supports_cdm = False + supports_ida = False supports_diversity_loss = False supports_real_data_fake = False def __init__(self, config): super().__init__(config) - if self.cdm_enabled: - raise ValueError("minimax_h3_t2av_dmd does not support CDM.") - if self.guidance_scale > 1.0: + model_name = self.model_config.get("name") + if model_name != "minimax_h3_t2av": + raise ValueError(f"minimax_h3_t2av_dmd requires model.name=minimax_h3_t2av, got {model_name!r}.") + if self.guidance_scale != 1.0: raise ValueError("MiniMax-H3 is guidance-distilled and has no unconditional branch; set training.teacher.guidance_scale to 1.0.") - h3 = self.training_config.get("minimax_h3", {}) - self.video_loss_weight = float(h3.get("video_loss_weight", 1.0)) - self.audio_loss_weight = float(h3.get("audio_loss_weight", 1.0)) - self.video_shift = float(h3.get("video_flow_shift", 12.0)) - self.audio_shift = float(h3.get("audio_flow_shift", 3.0)) - if self.video_shift <= 0 or self.audio_shift <= 0: - raise ValueError("MiniMax-H3 flow shifts must be positive.") - dtype_name = str(h3.get("latent_dtype", "fp32")).lower() - latent_dtypes = {"fp32": torch.float32, "float32": torch.float32, "bf16": torch.bfloat16} - if dtype_name not in latent_dtypes: - raise ValueError(f"Unsupported training.minimax_h3.latent_dtype={dtype_name!r}.") - self.latent_dtype = latent_dtypes[dtype_name] - self._layout_cache = {} - - # DMD's num_inference_steps means model evaluations. H3's published - # scheduler has one additional terminal-zero sigma grid point. - base = torch.linspace(1.0, 0.0, self.num_inference_steps + 1, dtype=torch.float32) - self.video_sigmas_cpu = _shift_sigma(base, self.video_shift) - self.audio_sigmas_cpu = _shift_sigma(base, self.audio_shift) - logger.info( - "[train] H3 DMD rollout evaluations={} video_shift={} audio_shift={}", - self.num_inference_steps, - self.video_shift, - self.audio_shift, - ) - - def _prepare_cached_condition(self, condition): - return self.model.prepare_text_condition(condition) - - def _encode_conditions(self, sample): - positive = sample["conditioning"].get("positive") - if positive is None: - raise KeyError("MiniMax-H3 DMD requires conditioning.positive. Build the prompt cache with build_minimax_h3_prompt_latents.py and use data.train.name=latent_dataset.") - with torch.no_grad(): - condition = self._prepare_cached_condition(positive) - return broadcast_sequence_parallel_value(condition), None - - def _latent_shape(self, sample): - del sample - dmd = self.dmd_config - height = int(dmd["height"]) - width = int(dmd["width"]) - num_frames = int(dmd["num_frames"]) - if height % 32 or width % 32: - raise ValueError(f"MiniMax-H3 height/width must be divisible by 32, 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 latent_height % patch_h or latent_width % patch_w: - raise ValueError(f"H3 latent canvas {latent_height}x{latent_width} is not divisible by patch {self.model.patch_size}.") - video_rows = latent_frames * (latent_height // patch_h) * (latent_width // patch_w) - video_dim = self.model.video_latent_channels * patch_t * patch_h * patch_w - audio_latents = audio_latent_num_frames(num_frames) - audio_rows = audio_latents * 2 - return { - "batch_size": 1, - "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_dim), - "audio_tokens": (1, audio_rows, self.model.audio_latent_channels), + if self.warp_denoising_step: + raise ValueError("MiniMax-H3 applies separate video/audio flow shifts in its model capability; set training.dmd.warp_denoising_step to false.") + + h3_config = self.training_config.get("minimax_h3", {}) + if not isinstance(h3_config, dict): + raise TypeError("training.minimax_h3 must be a mapping.") + dtype_name = str(h3_config.get("latent_dtype", "fp32")).lower() + dtype_aliases = { + "fp32": torch.float32, + "float32": torch.float32, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, } - - def sample_initial_latents(self, latent_shape): - video = torch.randn(latent_shape["video_tokens"], device=self.model.device, dtype=self.latent_dtype) - audio = torch.randn(latent_shape["audio_tokens"], device=self.model.device, dtype=self.latent_dtype) - return broadcast_sequence_parallel_value((video, audio)) - - def _layout(self, condition, latent_shape): - tags = condition["text_token_tags"] - if tags.ndim != 1: - raise ValueError(f"text_token_tags must be one-dimensional, got {tuple(tags.shape)}.") - if not bool((tags == 1).all()): - raise ValueError("The initial minimax_h3_t2av_dmd trainer only supports text-only conditions.") - key = ( - int(tags.numel()), - latent_shape["latent_frames"], - latent_shape["latent_height"], - latent_shape["latent_width"], - latent_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(), - latent_shape["latent_frames"], - latent_shape["latent_height"], - latent_shape["latent_width"], - latent_shape["audio_latents"], - self.model.patch_size, - ).to(self.model.device) - self._layout_cache[key] = layout - return layout - - def _predict_velocity(self, model, latents, sigmas, condition, latent_shape): - video, audio = latents - video_sigma, audio_sigma = sigmas - if video.shape[0] != 1 or audio.shape[0] != 1: - raise ValueError("MiniMax-H3 packed T2AV training currently requires batch_size=1.") - layout = self._layout(condition, latent_shape) - timesteps, timestep_indices = build_row_timesteps(layout, video_sigma, audio_sigma) - with model.transformer_forward_context(): - return model.denoiser_module()( - hidden_states=video, - audio_hidden_states=audio, - encoder_hidden_states=condition["prompt_embeds"], - timestep=timesteps.to(model.device), - timestep_indices=timestep_indices.to(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, - ) - - def run_back_simulation(self, condition, latent_shape, end_step_idx, grad_enabled, xt=None): - if xt is None: - xt = self.sample_initial_latents(latent_shape) - self.model.transformer.train() - x0 = None - for step_idx in range(end_step_idx + 1): - video_sigma = self.video_sigmas_cpu[step_idx].to(self.model.device) - audio_sigma = self.audio_sigmas_cpu[step_idx].to(self.model.device) - context = torch.enable_grad if grad_enabled and step_idx == end_step_idx else torch.no_grad - with context(): - velocity = self._predict_velocity(self.model, xt, (video_sigma, audio_sigma), condition, latent_shape) - x0 = ( - xt[0] + video_sigma * velocity[0], - xt[1] + audio_sigma * velocity[1], - ) - next_video_sigma = self.video_sigmas_cpu[step_idx + 1].to(self.model.device) - next_audio_sigma = self.audio_sigmas_cpu[step_idx + 1].to(self.model.device) - xt = ( - (xt[0].float() + (video_sigma - next_video_sigma) * velocity[0].float()).to(self.latent_dtype), - (xt[1].float() + (audio_sigma - next_audio_sigma) * velocity[1].float()).to(self.latent_dtype), - ) - return x0[0].to(self.latent_dtype), x0[1].to(self.latent_dtype) - - def _sample_renoise_sigmas(self): - low = float(self.dmd_config.get("renoise_sigma_min", 0.02)) - high = float(self.dmd_config.get("renoise_sigma_max", 0.98)) - if not 0.0 <= low < high <= 1.0: - raise ValueError(f"H3 renoise sigma range must satisfy 0 <= min < max <= 1, got [{low}, {high}].") - base = torch.empty((), device=self.model.device, dtype=torch.float32).uniform_(low, high) - return _shift_sigma(base, self.video_shift), _shift_sigma(base, self.audio_shift) - - @staticmethod - def _add_noise(latents, noises, sigmas): - return ( - ((1.0 - sigmas[0]) * latents[0].float() + sigmas[0] * noises[0]).to(latents[0].dtype), - ((1.0 - sigmas[1]) * latents[1].float() + sigmas[1] * noises[1]).to(latents[1].dtype), - ) - - def forward_loss(self, latent_shape, conditions, stage, current_iter=None, initial_noise=None): - del current_iter - condition, _ = conditions - end_step_idx = self._sample_synced_int(0, self.num_inference_steps) - generated = self.run_back_simulation( - condition, - latent_shape, - end_step_idx, - grad_enabled=stage != "fake", - xt=initial_noise, + if dtype_name not in dtype_aliases: + supported = ", ".join(sorted(dtype_aliases)) + raise ValueError(f"Unsupported training.minimax_h3.latent_dtype={dtype_name!r}; expected one of: {supported}.") + self.latent_dtype = dtype_aliases[dtype_name] + # Generic DMD uses running_dtype for rollout states and score sigmas. + # H3 defaults to fp32 latents even when prompt embeddings are bf16. + self.running_dtype = self.latent_dtype + + self.renoise_sigma_min = float(self.dmd_config.get("renoise_sigma_min", 0.02)) + self.renoise_sigma_max = float(self.dmd_config.get("renoise_sigma_max", 0.98)) + if not 0.0 <= self.renoise_sigma_min < self.renoise_sigma_max <= 1.0: + raise ValueError(f"MiniMax-H3 renoise sigma range must satisfy 0 <= min < max <= 1, got [{self.renoise_sigma_min}, {self.renoise_sigma_max}].") + logger.info( + "[train] MiniMax-H3 generic DMD latent_dtype={} renoise_sigma=[{}, {}]", + dtype_name, + self.renoise_sigma_min, + self.renoise_sigma_max, ) - sigmas = self._sample_renoise_sigmas() - noises = (torch.randn_like(generated[0], dtype=torch.float32), torch.randn_like(generated[1], dtype=torch.float32)) - renoised = self._add_noise(generated, noises, sigmas) - - if stage == "fake": - self.fake_model.transformer.train() - velocity_fake = self._predict_velocity(self.fake_model, renoised, sigmas, condition, latent_shape) - # H3 predicts clean-ward velocity d x(t) / dt = x0 - noise. - target = (generated[0].float() - noises[0], generated[1].float() - noises[1]) - return weighted_mse_pair( - velocity_fake, - target, - self.video_loss_weight, - self.audio_loss_weight, - ) - with torch.no_grad(): - self.fake_model.transformer.eval() - self.teacher_model.transformer.eval() - velocity_fake = self._predict_velocity(self.fake_model, renoised, sigmas, condition, latent_shape) - velocity_teacher = self._predict_velocity(self.teacher_model, renoised, sigmas, condition, latent_shape) - x0_fake = ( - renoised[0] + sigmas[0] * velocity_fake[0], - renoised[1] + sigmas[1] * velocity_fake[1], - ) - x0_teacher = ( - renoised[0] + sigmas[0] * velocity_teacher[0], - renoised[1] + sigmas[1] * velocity_teacher[1], - ) - loss = dmd_loss_pair( - generated, - x0_fake, - x0_teacher, - self.video_loss_weight, - self.audio_loss_weight, + def _sample_score_sigma( + self, + denoised_timestep_from, + denoised_timestep_to, + device, + dtype, + ): + del denoised_timestep_from, denoised_timestep_to + sigma = torch.empty((1,), device=device, dtype=torch.float32).uniform_( + self.renoise_sigma_min, + self.renoise_sigma_max, ) - return {"loss": loss, "dmd": loss.detach()} + return broadcast_sequence_parallel_value(sigma).to(dtype=dtype) diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index f18ac9e7c..d29883fa9 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -78,7 +78,7 @@ def merge(self, other_register): "longcat_image": "lightx2v_train.model_zoo.longcat_image.longcat_image", "ltx_t2av": "lightx2v_train.model_zoo.ltx.ltx_t2av", "ltx_t2av_ar": "lightx2v_train.model_zoo.ltx.ltx_t2av", - "minimax_h3_t2av": "lightx2v_train.model_zoo.minimax_h3_t2av", + "minimax_h3_t2av": "lightx2v_train.model_zoo.minimax_h3.minimax_h3_t2av", "qwen_image": "lightx2v_train.model_zoo.qwen_image.qwen_image", "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image.qwen_image_edit", "wan_t2v": "lightx2v_train.model_zoo.wan.wan_t2v", diff --git a/lightx2v_train/scripts/run_minimax_h3_t2av_dmd.sh b/lightx2v_train/scripts/run_minimax_h3_t2av_dmd.sh new file mode 100755 index 000000000..301780487 --- /dev/null +++ b/lightx2v_train/scripts/run_minimax_h3_t2av_dmd.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash + +set -euo pipefail + +H3_CONFIG_PATH="${H3_CONFIG_PATH:-configs/train/dmd/minimax_h3_t2av_dmd_lora.yaml}" +# The supplied config shards student, fake, and teacher across all eight GPUs. +# Override this when launching on a different node shape. +H3_NUM_PROCESSES="${H3_NUM_PROCESSES:-8}" + +torchrun \ + --standalone \ + --nproc_per_node="${H3_NUM_PROCESSES}" \ + train.py \ + --config "${H3_CONFIG_PATH}" From b01487ea88699971b35a0600c8243706004c6056 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 17 Aug 2026 16:09:18 +0000 Subject: [PATCH 08/34] update cm --- .../model_capabilities/training.py | 27 +- .../lightx2v_train/model_zoo/base.py | 87 +++--- .../model_zoo/capability_adapters/common.py | 40 ++- .../capability_adapters/consistency.py | 284 ++++++++++++++++++ .../flux2/capability_adapters/__init__.py | 4 + .../flux2_consistency_capability.py | 22 ++ .../model_zoo/flux2/flux2_dev.py | 7 +- .../model_zoo/flux2/flux2_klein.py | 7 +- .../capability_adapters/__init__.py | 5 + .../longcat_image_consistency_capability.py | 22 ++ .../model_zoo/longcat_image/longcat_image.py | 7 +- .../ltx/capability_adapters/__init__.py | 4 + .../ltx/capability_adapters/common.py | 101 +++++++ .../ltx_consistency_capability.py | 194 ++++++++++++ .../ltx_flow_matching_capability.py | 85 +----- .../lightx2v_train/model_zoo/ltx/ltx_t2av.py | 6 + .../capability_adapters/__init__.py | 2 + .../minimax_h3_consistency_capability.py | 235 +++++++++++++++ .../model_zoo/minimax_h3/minimax_h3_t2av.py | 8 +- .../capability_adapters/__init__.py | 5 + .../qwen_image_consistency_capability.py | 24 ++ .../model_zoo/qwen_image/qwen_image.py | 85 +----- .../model_zoo/qwen_image/qwen_image_edit.py | 7 +- .../wan/capability_adapters/__init__.py | 7 + .../wan_consistency_capability.py | 38 +++ .../model_zoo/wan/lingbot_video.py | 7 +- .../lightx2v_train/model_zoo/wan/wan_t2v.py | 6 + .../trainers/consistency/base.py | 1 + .../lightx2v_train/trainers/consistency/cm.py | 3 +- .../trainers/consistency/mean_flow.py | 18 +- .../trainers/consistency/pcm.py | 4 +- .../trainers/consistency/scm.py | 3 +- .../trainers/consistency/tcm.py | 3 +- .../trainers/consistency/trainer.py | 6 + 34 files changed, 1132 insertions(+), 232 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 9882f3379..4a4659804 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -1,6 +1,7 @@ from __future__ import annotations from abc import abstractmethod +from collections.abc import Collection from dataclasses import dataclass, field from typing import Any, Callable, Mapping @@ -54,36 +55,44 @@ def compute_loss( class ConsistencyCapability(ModelCapability): @abstractmethod - def configure(self, features) -> None: - pass + def configure(self, features: Collection[str]) -> None: + """Install model structure required by the selected objective.""" @abstractmethod def restore_trainable_auxiliary(self) -> None: - pass + """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): - pass + """Encode one training sample into the objective's clean state.""" @abstractmethod def encode_condition(self, batch): - pass + """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): - pass + """Run the model and convert its output to the requested parameterization.""" @abstractmethod def predict_log_variance(self, time): - pass + """Predict the learned loss log-variance required by sCM.""" @abstractmethod def set_frozen(self, training: bool = False) -> None: - pass + """Freeze the denoiser while selecting its train/eval execution mode.""" @abstractmethod def denoiser(self): - pass + """Return the module owned by this capability.""" class DistillationCapability(ModelCapability): diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 5e8dec550..2c4193a58 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -56,22 +56,6 @@ def load_components(self, transformer_only=False, reference_model=None): def denoiser_module(self): raise NotImplementedError(f"{self.__class__.__name__} must define denoiser_module().") - def configure_consistency_model(self, capabilities): - """Install optional architecture pieces required by an objective.""" - capabilities = frozenset(capabilities) - if capabilities: - names = ", ".join(sorted(capabilities)) - raise NotImplementedError(f"{self.__class__.__name__} does not support consistency capabilities: {names}.") - - def set_consistency_modules_trainable(self): - """Re-enable objective-specific modules after LoRA freezes the backbone.""" - - def consistency_auxiliary_parameter_names(self): - return () - - def predict_consistency_log_variance(self, time): - raise NotImplementedError(f"{self.__class__.__name__} does not implement a consistency log-variance head.") - def add_lora(self, rank, alpha, target_modules): lora_config = LoraConfig( r=rank, @@ -271,8 +255,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, auxiliary_state_dict = self._get_lora_and_auxiliary_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 @@ -284,15 +279,21 @@ def save_lora_weights(self, save_dir, adapter_name=None, weights_subdir=None): 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, "consistency_auxiliary.safetensors"), + 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): + 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): @@ -309,10 +310,10 @@ def _get_lora_and_auxiliary_state_dict_for_save(self, adapter_name=None): return {}, {} peft_state_dict = get_peft_model_state_dict(denoiser, state_dict=state_dict, **peft_kwargs) - auxiliary_names = set(self.consistency_auxiliary_parameter_names()) + auxiliary_names = set(auxiliary_parameter_names) missing = auxiliary_names - state_dict.keys() if missing: - raise RuntimeError(f"Consistency auxiliary parameters are missing from the model state: {sorted(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 @@ -331,44 +332,26 @@ 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 save_consistency_auxiliary_weights(self, save_dir): - names = set(self.consistency_auxiliary_parameter_names()) - if not names: - return - denoiser = self.denoiser_module() - if is_fsdp2_module(denoiser): - options = StateDictOptions( - full_state_dict=True, - cpu_offload=True, - ignore_frozen_params=False, - strict=False, - ) - state_dict, _ = get_state_dict(denoiser, (), options=options) - else: - state_dict = denoiser.state_dict() - if not is_main_process(): - return - missing = names - state_dict.keys() - if missing: - raise RuntimeError(f"Consistency auxiliary parameters are missing from the model state: {sorted(missing)}") - auxiliary = {name: state_dict[name].detach().cpu().contiguous() for name in names if name in state_dict} - if auxiliary: - save_file(auxiliary, os.path.join(save_dir, "consistency_auxiliary.safetensors")) - - def load_consistency_auxiliary_weights(self, checkpoint_dir): - names = set(self.consistency_auxiliary_parameter_names()) + 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, "consistency_auxiliary.safetensors") + path = os.path.join(checkpoint_dir, weights_name) if not os.path.exists(path): - raise RuntimeError(f"Consistency auxiliary weights were not found at {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 consistency auxiliary keys in {path}: {missing}") + raise RuntimeError(f"Missing auxiliary keys in {path}: {missing}") if unexpected: - logger.warning("Unexpected consistency auxiliary keys: {}", 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/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py index 235088366..1a3b9c7c7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from collections.abc import Collection from typing import Any, Mapping import torch @@ -110,9 +111,21 @@ def state_module(self): class CommonCheckpointCapability(BoundCapability, CheckpointCapability): + _CONSISTENCY_AUXILIARY_WEIGHTS_NAME = "consistency_auxiliary.safetensors" + + def _consistency_auxiliary_parameter_names(self) -> tuple[str, ...]: + capabilities = self.model.ensure_capabilities() + if not capabilities.supports(ConsistencyCapability): + return () + return capabilities.require(ConsistencyCapability).auxiliary_parameter_names() + def save_weights(self, save_dir, train_type) -> None: if train_type == "lora": - self.model.save_lora_weights(save_dir) + self.model.save_lora_weights( + save_dir, + auxiliary_parameter_names=self._consistency_auxiliary_parameter_names(), + auxiliary_weights_name=self._CONSISTENCY_AUXILIARY_WEIGHTS_NAME, + ) elif is_main_process(): torch.save( self.model.denoiser_module().state_dict(), @@ -122,7 +135,11 @@ def save_weights(self, save_dir, train_type) -> None: 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_consistency_auxiliary_weights(save_dir) + self.model.load_auxiliary_weights( + save_dir, + self._consistency_auxiliary_parameter_names(), + weights_name=self._CONSISTENCY_AUXILIARY_WEIGHTS_NAME, + ) return path = os.path.join(save_dir, "model_state.pt") if not os.path.exists(path): @@ -180,11 +197,17 @@ def compute_loss( class GenericConsistencyCapability(BoundCapability, ConsistencyCapability): - def configure(self, features) -> None: - self.model.configure_consistency_model(features) + 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: - self.model.set_consistency_modules_trainable() + pass + + def auxiliary_parameter_names(self) -> tuple[str, ...]: + return () def encode_latent(self, batch): return _require_singleton_tensor( @@ -195,6 +218,10 @@ def encode_latent(self, batch): def encode_condition(self, batch): return self.model.encode_condition(batch) + 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, @@ -211,7 +238,8 @@ def predict(self, request, path): ) def predict_log_variance(self, time): - return self.model.predict_consistency_log_variance(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() diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py new file mode 100644 index 000000000..34ec515ad --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.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 GenericConsistencyCapability + + +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 TimeConditionedConsistencyCapability(GenericConsistencyCapability): + """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/flux2/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py index 175635986..8150f7e40 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -1 +1,5 @@ """Capabilities implemented specifically by Flux2 models.""" + +from .flux2_consistency_capability import Flux2ConsistencyCapability + +__all__ = ["Flux2ConsistencyCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py new file mode 100644 index 000000000..3e4a42123 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py @@ -0,0 +1,22 @@ +"""Consistency capability for Flux2 backbones.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + ProjectedTimeEmbeddingAdapter, + TimeConditionedConsistencyCapability, +) + + +class Flux2ConsistencyCapability(TimeConditionedConsistencyCapability): + """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/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index 2bd5da334..8a1dc5387 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -4,8 +4,9 @@ from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor -from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -33,6 +34,10 @@ def register_capabilities(self): DistillationCapability, GenericDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + Flux2ConsistencyCapability(self), + ) def load_components(self, transformer_only=False, reference_model=None): model_config = self.config["model"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 09137d679..04228aaf3 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -4,8 +4,9 @@ from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor -from lightx2v_train.model_capabilities import DistillationCapability, DopsdCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, DopsdCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability from lightx2v_train.model_zoo.flux2.capability_adapters.flux2_dopsd_capability import ( Flux2DopsdCapability, ) @@ -36,6 +37,10 @@ def register_capabilities(self): DistillationCapability, GenericDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + Flux2ConsistencyCapability(self), + ) self.capabilities.register( DopsdCapability, Flux2DopsdCapability(self), 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..90927ea19 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/__init__.py @@ -0,0 +1,5 @@ +"""Capabilities implemented specifically by LongCat-Image.""" + +from .longcat_image_consistency_capability import LongCatImageConsistencyCapability + +__all__ = ["LongCatImageConsistencyCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py new file mode 100644 index 000000000..60549250c --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py @@ -0,0 +1,22 @@ +"""Consistency capability for LongCat-Image.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + ProjectedTimeEmbeddingAdapter, + TimeConditionedConsistencyCapability, +) + + +class LongCatImageConsistencyCapability(TimeConditionedConsistencyCapability): + """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/longcat_image.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py index d416206a7..35fd6dd66 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -6,8 +6,9 @@ from diffusers.models.transformers import LongCatImageTransformer2DModel from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids -from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.longcat_image.capability_adapters import LongCatImageConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -35,6 +36,10 @@ def register_capabilities(self): DistillationCapability, GenericDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + LongCatImageConsistencyCapability(self), + ) def load_components(self, transformer_only=False, reference_model=None): if transformer_only: diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py index e7c260825..bffaba5a2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py @@ -1 +1,5 @@ """Capabilities implemented specifically by LTX models.""" + +from .ltx_consistency_capability import LTXConsistencyCapability + +__all__ = ["LTXConsistencyCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py index 7b08f010d..5d1cd3784 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py @@ -2,8 +2,18 @@ from dataclasses import dataclass +import torch from torch import Tensor +from lightx2v_train.model_zoo.native.ltx2 import ( + AudioLatentShape, + AudioPatchifier, + SpatioTemporalScaleFactors, + VideoLatentPatchifier, + VideoLatentShape, + get_pixel_coords, +) + @dataclass(frozen=True) class LTXJointLatentShape: @@ -19,3 +29,94 @@ class LTXJointLatents: video: Tensor audio: Tensor shape: LTXJointLatentShape + + +class LTXLatentCodec: + """Validate, patchify, position, and flatten joint LTX latent tensors.""" + + def __init__(self, model, *, default_fps: float = 24.0) -> None: + self.model = model + self.default_fps = float(default_fps) + self.video_patchifier = VideoLatentPatchifier(patch_size=1) + self.audio_patchifier = AudioPatchifier(patch_size=1) + + def prepare(self, inputs, dtype): + video_data = inputs.get("video_latents") + audio_data = inputs.get("audio_latents") + if video_data is None or audio_data is None: + raise KeyError("LTX training expects inputs.video_latents and inputs.audio_latents.") + video_latents = self.video_latents(video_data, dtype) + audio_latents = self.audio_latents(audio_data, dtype) + video_tokens = self.video_patchifier.patchify(video_latents) + audio_tokens = self.audio_patchifier.patchify(audio_latents) + return video_data, audio_data, video_latents, audio_latents, video_tokens, audio_tokens + + def positions(self, video_data, video_latents, audio_tokens): + video_positions = self.video_positions(video_data, video_latents) + audio_positions = self.audio_patchifier.get_patch_grid_bounds( + output_shape=AudioLatentShape( + frames=audio_tokens.shape[1], + mel_bins=16, + batch=1, + channels=8, + ), + device=audio_tokens.device, + ) + return video_positions, audio_positions + + def flatten(self, video_tokens: Tensor, audio_tokens: Tensor) -> Tensor: + if video_tokens.shape[0] != 1 or audio_tokens.shape[0] != 1: + raise ValueError("LTX consistency training only supports physical batch size 1.") + return torch.cat((video_tokens.flatten(1), audio_tokens.flatten(1)), dim=1) + + @staticmethod + def unflatten(flat: Tensor, video_shape, audio_shape) -> tuple[Tensor, Tensor]: + video_numel = int(torch.Size(video_shape).numel()) + audio_numel = int(torch.Size(audio_shape).numel()) + if flat.shape != (1, video_numel + audio_numel): + raise ValueError(f"Flattened LTX latent must have shape {(1, video_numel + audio_numel)}, got {tuple(flat.shape)}.") + video, audio = flat.split((video_numel, audio_numel), dim=1) + return video.reshape(video_shape), audio.reshape(audio_shape) + + def video_latents(self, data, dtype): + latents = data["latents"].to(device=self.model.device, dtype=dtype) + if latents.ndim == 4: + latents = latents.unsqueeze(0) + if latents.ndim != 5: + raise ValueError(f"LTX video latents must have shape [B,C,F,H,W] or [C,F,H,W], got {tuple(latents.shape)}.") + if latents.shape[0] != 1: + raise ValueError("LTX training only supports physical batch size 1.") + return latents.contiguous() + + def audio_latents(self, data, dtype): + latents = data["latents"].to(device=self.model.device, dtype=dtype) + if latents.ndim == 3: + latents = latents.unsqueeze(0) + if latents.ndim != 4: + raise ValueError(f"LTX audio latents must have shape [B,C,T,F] or [C,T,F], got {tuple(latents.shape)}.") + if latents.shape[0] != 1: + raise ValueError("LTX training only supports physical batch size 1.") + return latents.contiguous() + + def video_positions(self, data, latents): + frames = int(data["num_frames"][0].item()) if "num_frames" in data else latents.shape[2] + height = int(data["height"][0].item()) if "height" in data else latents.shape[3] + width = int(data["width"][0].item()) if "width" in data else latents.shape[4] + fps = float(data["fps"][0].item()) if "fps" in data else self.default_fps + latent_coords = self.video_patchifier.get_patch_grid_bounds( + output_shape=VideoLatentShape( + frames=frames, + height=height, + width=width, + batch=1, + channels=128, + ), + device=latents.device, + ) + positions = get_pixel_coords( + latent_coords=latent_coords, + scale_factors=SpatioTemporalScaleFactors.default(), + causal_fix=True, + ).float() + positions[:, 0, ...] /= fps + return positions diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py new file mode 100644 index 000000000..403c37a74 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py @@ -0,0 +1,194 @@ +"""Consistency training for LTX's joint video/audio latent state.""" + +from __future__ import annotations + +import copy + +import torch +from torch import Tensor, nn + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + TimeConditionedConsistencyCapability, + TimeEmbeddingAdapter, +) +from lightx2v_train.model_zoo.native.ltx2 import Modality + +from .common import LTXLatentCodec + + +class LTXTimeEmbeddingAdapter(TimeEmbeddingAdapter): + """Condition both LTX modalities on a shared MeanFlow endpoint.""" + + def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: + return denoiser.adaln_single, denoiser.audio_adaln_single + + def base_embedder(self, denoiser: nn.Module) -> nn.Module: + return denoiser.adaln_single + + def build_endpoint_embedder(self, denoiser: nn.Module) -> nn.Module: + return nn.ModuleDict( + { + "video": copy.deepcopy(denoiser.adaln_single), + "audio": copy.deepcopy(denoiser.audio_adaln_single), + } + ) + + def embedding_dimension(self, denoiser: nn.Module) -> int: + return int(denoiser.inner_dim) + + def encode( + self, + denoiser: nn.Module, + embedder: nn.Module, + time: Tensor, + ) -> Tensor: + parameter = next(embedder.parameters()) + scaled_time = time.to(device=parameter.device) * float(denoiser.timestep_scale_multiplier) + _, embedding = embedder( + scaled_time.flatten(), + hidden_dtype=parameter.dtype, + ) + return embedding + + def add_endpoint_embedding( + self, + denoiser: nn.Module, + endpoint_embedder: nn.ModuleDict, + endpoint_time: Tensor, + hook_module: nn.Module, + output, + ): + if not isinstance(output, tuple) or len(output) != 2 or not all(torch.is_tensor(value) for value in output): + raise TypeError("LTX AdaLN endpoint conditioning expects a pair of tensor outputs.") + if hook_module is denoiser.adaln_single: + branch = endpoint_embedder["video"] + elif hook_module is denoiser.audio_adaln_single: + branch = endpoint_embedder["audio"] + else: + raise RuntimeError("Received an endpoint hook from an unknown LTX AdaLN module.") + + parameter = next(branch.parameters()) + scaled_time = endpoint_time.to(device=parameter.device) * float(denoiser.timestep_scale_multiplier) + endpoint_outputs = branch( + scaled_time.flatten(), + hidden_dtype=parameter.dtype, + ) + return tuple(self._add_broadcast(value, endpoint) for value, endpoint in zip(output, endpoint_outputs)) + + @staticmethod + def _add_broadcast(value: Tensor, endpoint: Tensor) -> Tensor: + while endpoint.ndim < value.ndim: + endpoint = endpoint.unsqueeze(-2) + return value + endpoint.to(device=value.device, dtype=value.dtype) + + +class LTXConsistencyCapability(TimeConditionedConsistencyCapability): + """Expose joint LTX audio/video denoising through the tensor CM boundary.""" + + def __init__(self, model) -> None: + super().__init__(model, LTXTimeEmbeddingAdapter()) + config = model.config["training"].get("ltx2", {}) + self.latent_codec = LTXLatentCodec( + model, + default_fps=float(config.get("default_fps", 24.0)), + ) + + def encode_latent(self, batch): + _, _, _, _, video_tokens, audio_tokens = self.latent_codec.prepare( + batch["inputs"], + self.model.running_dtype, + ) + return self.latent_codec.flatten(video_tokens, audio_tokens) + + def encode_condition(self, batch): + conditioning = batch["conditioning"] + positive = conditioning.get("positive") + if positive is None: + video_context, audio_context, context_mask = self.model.encode_prompt_condition(conditioning.get("prompt", "")) + else: + video_context, audio_context, context_mask = self.model.prepare_text_condition(positive) + + video_data, _, video_latents, _, video_tokens, audio_tokens = self.latent_codec.prepare( + batch["inputs"], + self.model.running_dtype, + ) + video_positions, audio_positions = self.latent_codec.positions( + video_data, + video_latents, + audio_tokens, + ) + return { + "video_context": video_context, + "audio_context": audio_context, + "context_mask": context_mask, + "video_positions": video_positions, + "audio_positions": audio_positions, + "video_shape": tuple(video_tokens.shape), + "audio_shape": tuple(audio_tokens.shape), + } + + def sampling_latent_hw(self, batch, clean) -> tuple[int, int]: + del clean + _, _, video_latents, _, _, _ = self.latent_codec.prepare( + batch["inputs"], + self.model.running_dtype, + ) + return int(video_latents.shape[-2]), int(video_latents.shape[-1]) + + def predict(self, request, path): + with self._configured_prediction(request) as prepared_request: + if prepared_request.model_kwargs: + names = ", ".join(sorted(prepared_request.model_kwargs)) + raise TypeError(f"Unsupported LTX consistency model arguments: {names}.") + prediction = self._predict_velocity( + prepared_request.sample, + prepared_request.time, + prepared_request.condition, + ) + return path.convert_prediction( + prepared_request.sample, + prediction, + prepared_request.time, + source_type=self.model.denoiser_prediction_type(), + target_type=prepared_request.prediction_type, + ) + + def _predict_velocity(self, sample, time, condition): + video_tokens, audio_tokens = self.latent_codec.unflatten( + sample, + condition["video_shape"], + condition["audio_shape"], + ) + sigma = time.to(device=sample.device, dtype=sample.dtype).reshape(-1) + if sigma.numel() != 1: + raise ValueError(f"LTX consistency training requires one shared time value, got {tuple(time.shape)}.") + video_timesteps = sigma.view(1, 1).expand(1, video_tokens.shape[1]) + audio_timesteps = sigma.view(1, 1).expand(1, audio_tokens.shape[1]) + + video = Modality( + enabled=True, + latent=video_tokens, + sigma=sigma, + timesteps=video_timesteps, + positions=condition["video_positions"], + context=condition["video_context"], + context_mask=condition["context_mask"], + ) + audio = Modality( + enabled=True, + latent=audio_tokens, + sigma=sigma, + timesteps=audio_timesteps, + positions=condition["audio_positions"], + context=condition["audio_context"], + context_mask=condition["context_mask"], + ) + with self.model.transformer_forward_context(): + video_prediction, audio_prediction = self.denoiser()( + video=video, + audio=audio, + perturbations=None, + ) + video_prediction = video_prediction[:, -video_tokens.shape[1] :] + audio_prediction = audio_prediction[:, -audio_tokens.shape[1] :] + return self.latent_codec.flatten(video_prediction, audio_prediction) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py index 3d3ff3a59..b1d684a6e 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py @@ -13,15 +13,9 @@ LossResult, SFTStepContext, ) -from lightx2v_train.model_zoo.native.ltx2 import ( - AudioLatentShape, - AudioPatchifier, - Modality, - SpatioTemporalScaleFactors, - VideoLatentPatchifier, - VideoLatentShape, - get_pixel_coords, -) +from lightx2v_train.model_zoo.native.ltx2 import Modality + +from .common import LTXLatentCodec class LTXFlowMatchingCapability(BoundCapability, FlowMatchingSFTCapability): @@ -33,8 +27,7 @@ def __init__(self, model) -> None: self.video_loss_weight = float(config.get("video_loss_weight", 1.0)) self.audio_loss_weight = float(config.get("audio_loss_weight", 1.0)) self.default_fps = float(config.get("default_fps", 24.0)) - self.video_patchifier = VideoLatentPatchifier(patch_size=1) - self.audio_patchifier = AudioPatchifier(patch_size=1) + self.latent_codec = LTXLatentCodec(model, default_fps=self.default_fps) def compute_loss( self, @@ -42,16 +35,13 @@ def compute_loss( context: SFTStepContext, ) -> LossResult: inputs = batch["inputs"] - video_data = inputs.get("video_latents") - audio_data = inputs.get("audio_latents") conditions = batch["conditioning"].get("positive") - if video_data is None or audio_data is None or conditions is None: - raise KeyError("LTX flow matching expects inputs.video_latents, inputs.audio_latents, and conditioning.positive.") - - video_latents = self._video_latents(video_data, context.running_dtype) - audio_latents = self._audio_latents(audio_data, context.running_dtype) - video_tokens = self.video_patchifier.patchify(video_latents) - audio_tokens = self.audio_patchifier.patchify(audio_latents) + if conditions is None: + raise KeyError("LTX flow matching expects conditioning.positive.") + video_data, _, video_latents, _, video_tokens, audio_tokens = self.latent_codec.prepare( + inputs, + context.running_dtype, + ) video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) video_noisy = self._initialize_noisy_target( @@ -66,18 +56,10 @@ def compute_loss( video_input, video_target, video_timesteps, video_mask, video_sigmas = video_noisy audio_input, audio_target, audio_timesteps, audio_mask, audio_sigmas = audio_noisy - video_positions = self._video_positions( + video_positions, audio_positions = self.latent_codec.positions( video_data, video_latents, - ) - audio_positions = self.audio_patchifier.get_patch_grid_bounds( - output_shape=AudioLatentShape( - frames=audio_tokens.shape[1], - mel_bins=16, - batch=1, - channels=8, - ), - device=audio_tokens.device, + audio_tokens, ) video = Modality( enabled=True, @@ -124,26 +106,6 @@ def compute_loss( }, ) - def _video_latents(self, data, dtype): - latents = data["latents"].to(device=self.model.device, dtype=dtype) - if latents.ndim == 4: - latents = latents.unsqueeze(0) - if latents.ndim != 5: - raise ValueError(f"LTX video latents must have shape [B,C,F,H,W] or [C,F,H,W], got {tuple(latents.shape)}.") - if latents.shape[0] != 1: - raise ValueError("LTX flow matching only supports physical batch size 1.") - return latents.contiguous() - - def _audio_latents(self, data, dtype): - latents = data["latents"].to(device=self.model.device, dtype=dtype) - if latents.ndim == 3: - latents = latents.unsqueeze(0) - if latents.ndim != 4: - raise ValueError(f"LTX audio latents must have shape [B,C,T,F] or [C,T,F], got {tuple(latents.shape)}.") - if latents.shape[0] != 1: - raise ValueError("LTX flow matching only supports physical batch size 1.") - return latents.contiguous() - @staticmethod def _initialize_noisy_target(latents: Tensor, scheduler, sigmas=None): _, sequence_length, _ = latents.shape @@ -172,29 +134,6 @@ def _initialize_noisy_target(latents: Tensor, scheduler, sigmas=None): ) return noisy, target, timesteps, mask, sigmas - def _video_positions(self, data, latents): - frames = int(data["num_frames"][0].item()) if "num_frames" in data else latents.shape[2] - height = int(data["height"][0].item()) if "height" in data else latents.shape[3] - width = int(data["width"][0].item()) if "width" in data else latents.shape[4] - fps = float(data["fps"][0].item()) if "fps" in data else self.default_fps - latent_coords = self.video_patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape( - frames=frames, - height=height, - width=width, - batch=1, - channels=128, - ), - device=latents.device, - ) - positions = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=SpatioTemporalScaleFactors.default(), - causal_fix=True, - ).float() - positions[:, 0, ...] /= fps - return positions - @staticmethod def _masked_mse(prediction, target, mask): target_length = target.shape[1] diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py index 57b59c7d7..064499d63 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py @@ -11,6 +11,7 @@ from lightx2v_train.model_capabilities import ( AutoregressiveDistillationCapability, + ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability, TeacherForcingCapability, @@ -18,6 +19,7 @@ from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_autoregressive_distillation_capability import ( LTXAutoregressiveDistillationCapability, ) +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_consistency_capability import LTXConsistencyCapability from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_distillation_capability import LTXDistillationCapability from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_flow_matching_capability import LTXFlowMatchingCapability from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_teacher_forcing_capability import LTXTeacherForcingCapability @@ -55,6 +57,10 @@ def register_capabilities(self): DistillationCapability, LTXDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + LTXConsistencyCapability(self), + ) if self.config["model"]["name"] == "ltx_t2av_ar": self.capabilities.register( AutoregressiveDistillationCapability, 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 index 2788ef6e3..194f1b44b 100644 --- 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 @@ -1,9 +1,11 @@ """Capability adapters for MiniMax-H3.""" from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape +from .minimax_h3_consistency_capability import MiniMaxH3ConsistencyCapability from .minimax_h3_distillation_capability import MiniMaxH3DistillationCapability __all__ = [ + "MiniMaxH3ConsistencyCapability", "MiniMaxH3DistillationCapability", "MiniMaxH3JointLatents", "MiniMaxH3LatentShape", diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py new file mode 100644 index 000000000..be05db679 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py @@ -0,0 +1,235 @@ +"""Consistency training for MiniMax-H3 joint video/audio latents.""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor, nn + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + TimeConditionedConsistencyCapability, + TimeEmbeddingAdapter, +) +from lightx2v_train.model_zoo.native.minimax_h3 import MiniMaxH3PackedSequence, build_packed_sequence + + +class MiniMaxH3TimeEmbeddingAdapter(TimeEmbeddingAdapter): + """Translate noise time ``sigma`` to H3's clean-ward ``1 - sigma`` time.""" + + def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: + return (denoiser.time_embedder,) + + def base_embedder(self, denoiser: nn.Module) -> nn.Module: + return denoiser.time_embedder + + def embedding_dimension(self, denoiser: nn.Module) -> int: + return int(denoiser.config.time_embed_dim) + + def encode( + self, + denoiser: nn.Module, + embedder: nn.Module, + time: Tensor, + ) -> Tensor: + parameter = next(embedder.parameters()) + cleanward_time = 1.0 - time.to(device=parameter.device).reshape(-1) + projected = denoiser.time_proj(cleanward_time) + return embedder(projected.to(device=parameter.device, dtype=parameter.dtype)) + + +class MiniMaxH3ConsistencyCapability(TimeConditionedConsistencyCapability): + """Expose H3's packed joint denoiser through the consistency tensor API.""" + + def __init__(self, model) -> None: + super().__init__(model, MiniMaxH3TimeEmbeddingAdapter()) + + def encode_latent(self, batch): + video_tokens, _, _, _ = self._video_tokens(batch["inputs"]) + audio_tokens = self._audio_tokens(batch["inputs"]) + return self._flatten(video_tokens, audio_tokens) + + def encode_condition(self, batch): + positive = batch["conditioning"].get("positive") + if positive is None: + raise KeyError("MiniMax-H3 consistency training requires conditioning.positive with cached prompt embeddings.") + text_condition = self.model.prepare_text_condition(positive) + tags = text_condition["text_token_tags"] + if not bool((tags == 1).all()): + raise ValueError("MiniMax-H3 consistency training currently supports text-only cached conditions.") + + video_tokens, latent_frames, latent_height, latent_width = self._video_tokens(batch["inputs"]) + audio_tokens = self._audio_tokens(batch["inputs"]) + if audio_tokens.shape[1] % 2: + raise ValueError(f"MiniMax-H3 audio token count must be even for stereo audio, got {audio_tokens.shape[1]}.") + layout = build_packed_sequence( + tags.detach().cpu(), + latent_frames, + latent_height, + latent_width, + audio_tokens.shape[1] // 2, + self.model.patch_size, + ).to(self.model.device) + if layout.video_indices.numel() != video_tokens.shape[1]: + raise ValueError(f"MiniMax-H3 layout has {layout.video_indices.numel()} video rows but the cache contains {video_tokens.shape[1]}.") + if layout.audio_indices.numel() != audio_tokens.shape[1]: + raise ValueError(f"MiniMax-H3 layout has {layout.audio_indices.numel()} audio rows but the cache contains {audio_tokens.shape[1]}.") + return { + **text_condition, + "sequence_length": layout.sequence_length, + "position_ids": layout.position_ids, + "token_tags": layout.token_tags, + "video_indices": layout.video_indices, + "audio_indices": layout.audio_indices, + "text_indices": layout.text_indices, + "video_shape": tuple(video_tokens.shape), + "audio_shape": tuple(audio_tokens.shape), + } + + def sampling_latent_hw(self, batch, clean) -> tuple[int, int]: + del clean + _, _, latent_height, latent_width = self._video_tokens(batch["inputs"]) + return latent_height, latent_width + + def predict(self, request, path): + with self._configured_prediction(request) as prepared_request: + if prepared_request.model_kwargs: + names = ", ".join(sorted(prepared_request.model_kwargs)) + raise TypeError(f"Unsupported MiniMax-H3 consistency model arguments: {names}.") + prediction = self._predict_velocity( + prepared_request.sample, + prepared_request.time, + prepared_request.condition, + ) + return path.convert_prediction( + prepared_request.sample, + prediction, + prepared_request.time, + source_type="velocity", + target_type=prepared_request.prediction_type, + ) + + def _predict_velocity(self, sample, sigma, condition): + video_tokens, audio_tokens = self._unflatten( + sample, + condition["video_shape"], + condition["audio_shape"], + ) + sigma = sigma.to(device=sample.device, dtype=torch.float32).reshape(-1) + if sigma.numel() != 1: + raise ValueError(f"MiniMax-H3 consistency training requires one shared time value, got {tuple(sigma.shape)}.") + layout = self._layout(condition) + # Keep this conversion differentiable for sCM/MeanFlow forward-mode JVP. + timesteps = 1.0 - sigma + timestep_indices = torch.zeros(layout.sequence_length, device=self.model.device, dtype=torch.long) + with self.model.transformer_forward_context(): + prediction = self.denoiser()( + hidden_states=video_tokens, + audio_hidden_states=audio_tokens, + 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(prediction, (tuple, list)) or len(prediction) < 2: + raise TypeError("MiniMax-H3 transformer must return video and audio velocity tensors.") + # H3 predicts x0-noise, opposite to the framework's noise-x0 RF velocity. + return self._flatten(-prediction[0], -prediction[1]) + + def _video_tokens(self, inputs): + value = inputs.get("video_latents") + if value is None: + raise KeyError("MiniMax-H3 consistency training requires inputs.video_latents.") + metadata = value if isinstance(value, dict) else {} + tensor = self._cached_tensor(value, "video_latents") + if tensor.ndim == 2: + tensor = tensor.unsqueeze(0) + if tensor.ndim == 3: + tokens = tensor.to(device=self.model.device, dtype=self.model.running_dtype) + latent_frames = metadata.get("latent_frames") + latent_height = metadata.get("latent_height") + latent_width = metadata.get("latent_width") + if latent_frames is None or latent_height is None or latent_width is None: + raise KeyError("Patchified MiniMax-H3 video latents require latent_frames, latent_height, and latent_width metadata.") + geometry = int(latent_frames), int(latent_height), int(latent_width) + elif tensor.ndim in {4, 5}: + if tensor.ndim == 4: + tensor = tensor.unsqueeze(0) + tensor = tensor.to(device=self.model.device, dtype=self.model.running_dtype) + tokens = self._patchify_video(tensor) + geometry = int(tensor.shape[2]), int(tensor.shape[3]), int(tensor.shape[4]) + 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) + if tokens.shape[0] != 1 or tokens.shape[-1] != expected_dimension: + raise ValueError(f"MiniMax-H3 video tokens must have shape [1,N,{expected_dimension}], got {tuple(tokens.shape)}.") + return tokens.contiguous(), *geometry + + def _audio_tokens(self, inputs): + value = inputs.get("audio_latents") + if value is None: + raise KeyError("MiniMax-H3 consistency training requires inputs.audio_latents.") + 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)}.") + return tensor.to(device=self.model.device, dtype=self.model.running_dtype).contiguous() + + @staticmethod + def _cached_tensor(value, name): + if isinstance(value, dict): + 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 + + def _patchify_video(self, video): + batch, channels, frames, height, width = video.shape + patch_t, patch_h, patch_w = self.model.patch_size + if frames % patch_t or height % patch_h or width % patch_w: + raise ValueError(f"MiniMax-H3 raw video latent shape {tuple(video.shape)} is not divisible by patch size {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 _flatten(video, audio): + return torch.cat((video.flatten(1), audio.flatten(1)), dim=1) + + @staticmethod + def _unflatten(flat, video_shape, audio_shape): + video_numel = math.prod(video_shape) + audio_numel = math.prod(audio_shape) + if flat.shape != (1, video_numel + audio_numel): + raise ValueError(f"Flattened MiniMax-H3 latent must have shape {(1, video_numel + audio_numel)}, got {tuple(flat.shape)}.") + video, audio = flat.split((video_numel, audio_numel), dim=1) + return video.reshape(video_shape), audio.reshape(audio_shape) + + @staticmethod + def _layout(condition): + return MiniMaxH3PackedSequence( + sequence_length=int(condition["sequence_length"]), + position_ids=condition["position_ids"], + token_tags=condition["token_tags"], + video_indices=condition["video_indices"], + audio_indices=condition["audio_indices"], + text_indices=condition["text_indices"], + ) 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 index 0b3d4d066..2e9465869 100644 --- 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 @@ -5,8 +5,8 @@ import torch from peft import LoraConfig, inject_adapter_in_model -from lightx2v_train.model_capabilities import DistillationCapability -from lightx2v_train.model_zoo.minimax_h3.capability_adapters import MiniMaxH3DistillationCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability +from lightx2v_train.model_zoo.minimax_h3.capability_adapters import MiniMaxH3ConsistencyCapability, MiniMaxH3DistillationCapability 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 @@ -26,6 +26,10 @@ def register_capabilities(self): DistillationCapability, MiniMaxH3DistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + MiniMaxH3ConsistencyCapability(self), + ) def load_components(self, transformer_only=False, reference_model=None): del transformer_only, reference_model diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py new file mode 100644 index 000000000..a80ead5fb --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py @@ -0,0 +1,5 @@ +"""Capability adapters for Qwen-Image models.""" + +from .qwen_image_consistency_capability import QwenImageConsistencyCapability + +__all__ = ["QwenImageConsistencyCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py new file mode 100644 index 000000000..e2a6ac018 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py @@ -0,0 +1,24 @@ +"""Consistency-model capability for Qwen-Image backbones.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + ProjectedTimeEmbeddingAdapter, + TimeConditionedConsistencyCapability, +) + + +class QwenImageConsistencyCapability(TimeConditionedConsistencyCapability): + """Bind the generic consistency extensions to Qwen's time embedding.""" + + def __init__(self, model) -> None: + super().__init__( + model, + ProjectedTimeEmbeddingAdapter( + hook_module_path="time_text_embed", + projection_module_path="time_text_embed.time_proj", + embedding_module_path="time_text_embed.timestep_embedder", + embedding_dimension_path="inner_dim", + ), + # Keep existing Qwen consistency checkpoints byte-for-byte compatible. + endpoint_module_name="r_timestep_embedder", + log_variance_module_name="logvar_linear", + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index c13c0443d..ef2acfd6e 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -1,13 +1,12 @@ -import copy from dataclasses import dataclass import torch from diffusers import AutoencoderKLQwenImage, QwenImagePipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor -from torch import nn from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability -from lightx2v_train.model_zoo.capability_adapters.common import GenericConsistencyCapability, GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -42,7 +41,7 @@ def register_capabilities(self): ) self.capabilities.register( ConsistencyCapability, - GenericConsistencyCapability(self), + QwenImageConsistencyCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): @@ -81,57 +80,6 @@ def denoiser_prediction_type(self): # the corresponding velocity noise - x_0. return "velocity" - def configure_consistency_model(self, capabilities): - capabilities = frozenset(capabilities) - unsupported = capabilities - {"endpoint_time", "log_variance"} - if unsupported: - names = ", ".join(sorted(unsupported)) - raise NotImplementedError(f"QwenImageModel does not support consistency capabilities: {names}.") - - transformer = self.transformer - if "log_variance" in capabilities and not hasattr(transformer, "logvar_linear"): - transformer.logvar_linear = nn.Linear(transformer.inner_dim, 1).to( - device=self.device, - dtype=self.running_dtype, - ) - - if "endpoint_time" in capabilities: - if not hasattr(transformer, "r_timestep_embedder"): - transformer.r_timestep_embedder = copy.deepcopy(transformer.time_text_embed.timestep_embedder) - self._consistency_endpoint_time = None - if "endpoint_time" in capabilities and not hasattr(self, "_consistency_endpoint_hook"): - self._consistency_endpoint_hook = transformer.time_text_embed.register_forward_hook(self._add_endpoint_time_embedding) - self._consistency_capabilities = getattr(self, "_consistency_capabilities", frozenset()) | capabilities - - def _add_endpoint_time_embedding(self, time_embedder, inputs, output): - endpoint_time = getattr(self, "_consistency_endpoint_time", None) - if endpoint_time is None: - return output - projected = time_embedder.time_proj(endpoint_time) - projected = projected.to(device=output.device, dtype=output.dtype) - endpoint_embedding = self.transformer.r_timestep_embedder(projected) - return output + endpoint_embedding.to(dtype=output.dtype) - - def set_consistency_modules_trainable(self): - transformer = self.transformer - for name in ("logvar_linear", "r_timestep_embedder"): - module = getattr(transformer, name, None) - if module is not None: - module.requires_grad_(True) - - def consistency_auxiliary_parameter_names(self): - prefixes = ("logvar_linear.", "r_timestep_embedder.") - return tuple(name for name, _ in self.transformer.named_parameters() if name.startswith(prefixes)) - - def predict_consistency_log_variance(self, time): - if not hasattr(self.transformer, "logvar_linear"): - raise RuntimeError("The Qwen consistency log-variance head has not been configured.") - time_embedder = self.transformer.time_text_embed - projected = time_embedder.time_proj(time) - dtype = self.transformer.logvar_linear.weight.dtype - embedding = time_embedder.timestep_embedder(projected.to(device=self.device, dtype=dtype)) - return self.transformer.logvar_linear(embedding) - def fsdp2_shard_plan(self, fsdp_config): reshard_config = fsdp_config["reshard_after_forward"] return [ @@ -182,23 +130,16 @@ def prepare_denoiser_input(self, noisy_latent, condition=None): width=w, ) - def denoise(self, denoiser_input, timestep_or_sigma, condition, endpoint_time=None): - if endpoint_time is not None and not hasattr(self.transformer, "r_timestep_embedder"): - raise RuntimeError("endpoint_time was provided before Qwen endpoint conditioning was configured.") - previous_endpoint = getattr(self, "_consistency_endpoint_time", None) - self._consistency_endpoint_time = endpoint_time - try: - return self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, # timestep_or_sigma is in [0, 1] not [0, 1000] - guidance=None, - encoder_hidden_states_mask=condition["prompt_embed_mask"], - encoder_hidden_states=condition["prompt_embed"], - img_shapes=denoiser_input.img_shapes, - return_dict=False, - )[0] - finally: - self._consistency_endpoint_time = previous_endpoint + def denoise(self, denoiser_input, timestep_or_sigma, condition): + return self.transformer( + hidden_states=denoiser_input.hidden_states, + timestep=timestep_or_sigma, # timestep_or_sigma is in [0, 1] not [0, 1000] + guidance=None, + encoder_hidden_states_mask=condition["prompt_embed_mask"], + encoder_hidden_states=condition["prompt_embed"], + img_shapes=denoiser_input.img_shapes, + return_dict=False, + )[0] def postprocess_denoiser_output(self, prediction, denoiser_input): return QwenImagePipeline._unpack_latents( diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index 00388f80b..1f95206cf 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -6,8 +6,9 @@ from diffusers import AutoencoderKLQwenImage, QwenImageEditPlusPipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor -from lightx2v_train.model_capabilities import FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability +from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -50,6 +51,10 @@ def register_capabilities(self): FlowMatchingSFTCapability, GenericFlowMatchingCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + QwenImageConsistencyCapability(self), + ) def load_components(self, transformer_only=False, reference_model=None): if transformer_only: diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py index bc9c89a4e..8b32bda2b 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py @@ -1 +1,8 @@ """Capabilities implemented specifically by Wan-family models.""" + +from .wan_consistency_capability import LingBotConsistencyCapability, WanConsistencyCapability + +__all__ = [ + "LingBotConsistencyCapability", + "WanConsistencyCapability", +] diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py new file mode 100644 index 000000000..b379502f3 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py @@ -0,0 +1,38 @@ +"""Consistency capabilities for Wan-family video backbones.""" + +from lightx2v_train.model_zoo.capability_adapters.consistency import ( + ProjectedTimeEmbeddingAdapter, + SinusoidalTimeEmbeddingAdapter, + TimeConditionedConsistencyCapability, +) + + +class WanConsistencyCapability(TimeConditionedConsistencyCapability): + """Bind generic consistency extensions to Wan's sinusoidal time MLP.""" + + def __init__(self, model) -> None: + super().__init__( + model, + SinusoidalTimeEmbeddingAdapter( + embedding_module_path="time_embedding", + embedding_dimension_path="dim", + frequency_dimension_path="freq_dim", + time_scale=float(model.num_train_timesteps), + ), + ) + + +class LingBotConsistencyCapability(TimeConditionedConsistencyCapability): + """Bind generic consistency extensions to LingBot's time embedding.""" + + def __init__(self, model) -> None: + super().__init__( + model, + ProjectedTimeEmbeddingAdapter( + hook_module_path="time_embedder", + projection_module_path="time_proj", + embedding_module_path="time_embedder", + embedding_dimension_path="config.hidden_size", + time_scale=1000.0, + ), + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index 174d22b54..7dead91ce 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -9,8 +9,9 @@ from peft.utils import set_peft_model_state_dict from safetensors.torch import load_file -from lightx2v_train.model_capabilities import DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability +from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_capability import LingBotConsistencyCapability from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( LingBotDistillationCapability, ) @@ -55,6 +56,10 @@ def register_capabilities(self): DistillationCapability, LingBotDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + LingBotConsistencyCapability(self), + ) vae_scale_factor_temporal = 4 vae_scale_factor_spatial = 8 diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py index a0f328579..9d1623011 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py @@ -13,6 +13,7 @@ from lightx2v_train.model_capabilities import ( AutoregressiveDistillationCapability, + ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability, TeacherForcingCapability, @@ -21,6 +22,7 @@ from lightx2v_train.model_zoo.wan.capability_adapters.wan_autoregressive_distillation_capability import ( WanAutoregressiveDistillationCapability, ) +from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_capability import WanConsistencyCapability from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( WanDistillationCapability, ) @@ -60,6 +62,10 @@ def register_capabilities(self): DistillationCapability, WanDistillationCapability(self), ) + self.capabilities.register( + ConsistencyCapability, + WanConsistencyCapability(self), + ) if self.use_causal_transformer: self.capabilities.register( AutoregressiveDistillationCapability, diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/base.py b/lightx2v_train/lightx2v_train/trainers/consistency/base.py index 3b4d49cd4..38fa4a111 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/base.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/base.py @@ -118,6 +118,7 @@ class ConsistencyBatch: class ConsistencyStepContext: iteration: int global_batch_size: int + latent_hw: tuple[int, int] @dataclass diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/cm.py b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py index 33cdf4edc..57da67eae 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/cm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/cm.py @@ -162,8 +162,7 @@ def sample_training_state( context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: require_singleton_clean(clean) - latent_hw = (clean.shape[-2], clean.shape[-1]) - sampled_t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device) + sampled_t = scheduler.sample_timestep_or_sigma(latent_hw=context.latent_hw).to(clean.device) pair = self.time_pair_sampler.sample(sampled_t, context) return { "noise": torch.randn_like(clean), diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py index dea5a345e..093c8bf91 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/mean_flow.py @@ -23,10 +23,21 @@ def _where_condition(mask: Tensor, positive: Any, negative: Any): if torch.is_tensor(positive): + if not torch.is_tensor(negative): + raise TypeError(f"MeanFlow condition branches have incompatible types: Tensor and {type(negative).__name__}.") return torch.where(expand_time(mask, positive.ndim), positive, negative) if isinstance(positive, Mapping): + if not isinstance(negative, Mapping) or positive.keys() != negative.keys(): + raise ValueError("MeanFlow condition mappings must have identical keys.") return {key: _where_condition(mask, value, negative[key]) for key, value in positive.items()} - raise TypeError(f"MeanFlow condition dropout does not support {type(positive).__name__}.") + if isinstance(positive, (list, tuple)): + if not isinstance(negative, type(positive)) or len(positive) != len(negative): + raise ValueError("MeanFlow condition sequences must have identical types and lengths.") + values = [_where_condition(mask, value, negative[index]) for index, value in enumerate(positive)] + return type(positive)(values) + if positive == negative: + return positive + raise ValueError(f"MeanFlow cannot mix non-tensor condition metadata with different values: {positive!r} and {negative!r}.") @CONSISTENCY_OBJECTIVE_REGISTER("mean_flow") @@ -59,9 +70,8 @@ def sample_training_state( context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: require_singleton_clean(clean) - latent_hw = (clean.shape[-2], clean.shape[-1]) - first = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() - second = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() + first = scheduler.sample_timestep_or_sigma(latent_hw=context.latent_hw).to(clean.device).float() + second = scheduler.sample_timestep_or_sigma(latent_hw=context.latent_hw).to(clean.device).float() t = torch.maximum(first, second) random_r = torch.minimum(first, second) use_random_r = torch.rand(1, device=clean.device) < self.config.random_endpoint_probability diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py index 3bcbebc41..345192fe4 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -133,12 +133,10 @@ def sample_training_state( scheduler, context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: - del context require_singleton_clean(clean) - latent_hw = (clean.shape[-2], clean.shape[-1]) time = self.time_grid.sample( scheduler, - latent_hw=latent_hw, + latent_hw=context.latent_hw, device=clean.device, ) return { diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/scm.py b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py index cd9e1e69b..4ee0db7ac 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/scm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/scm.py @@ -49,8 +49,7 @@ def sample_training_state( context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: require_singleton_clean(clean) - latent_hw = (clean.shape[-2], clean.shape[-1]) - t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() + t = scheduler.sample_timestep_or_sigma(latent_hw=context.latent_hw).to(clean.device).float() return { "noise": torch.randn_like(clean) * self.config.sigma_data, "t": t, diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py index 66a4b8fcd..47daaa860 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/tcm.py @@ -90,8 +90,7 @@ def sample_training_state( context: ConsistencyStepContext, ) -> Mapping[str, Tensor]: require_singleton_clean(clean) - latent_hw = (clean.shape[-2], clean.shape[-1]) - t = scheduler.sample_timestep_or_sigma(latent_hw=latent_hw).to(clean.device).float() + t = scheduler.sample_timestep_or_sigma(latent_hw=context.latent_hw).to(clean.device).float() t = t.clamp_min(self.tcm_config.transition_time + self.config.time_pair.safety_epsilon) boundary_mask = torch.rand(1, device=clean.device) < self.tcm_config.boundary_probability diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py index 8151ba3af..7ed329dc0 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py @@ -215,6 +215,7 @@ def compute_loss_on_sample(self, sample): context = ConsistencyStepContext( iteration=int(getattr(self, "current_train_iteration", 0)), global_batch_size=(clean.shape[0] * get_data_parallel_world_size() * self.gradient_accumulation_iters), + latent_hw=self.consistency.sampling_latent_hw(sample, clean), ) training_state = self.objective.sample_training_state( clean, @@ -256,5 +257,10 @@ def _encode_negative_condition(self, sample): negative_sample = dict(sample) negative_sample["conditioning"] = dict(conditioning) + cached_negative = negative_sample["conditioning"].get("negative") + if cached_negative is None: + negative_sample["conditioning"].pop("positive", None) + else: + negative_sample["conditioning"]["positive"] = cached_negative negative_sample["conditioning"]["prompt"] = encoded_prompt return self.consistency.encode_condition(negative_sample) From f7703b6e4ddcf8a5f12d575d776880913773048d Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 17 Aug 2026 16:45:43 +0000 Subject: [PATCH 09/34] update --- .../lightx2v_train/model_zoo/base.py | 6 -- .../model_zoo/capability_adapters/__init__.py | 4 ++ .../model_zoo/capability_adapters/common.py | 23 ++++-- .../capability_adapters/latent_geometry.py | 71 +++++++++++++++++++ .../model_zoo/flux2/flux2_dev.py | 18 ++--- .../model_zoo/flux2/flux2_klein.py | 19 +++-- .../model_zoo/longcat_image/longcat_image.py | 17 +++-- .../model_zoo/qwen_image/qwen_image.py | 18 +++-- .../model_zoo/wan/lingbot_video.py | 7 +- 9 files changed, 127 insertions(+), 56 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/capability_adapters/latent_geometry.py diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 2c4193a58..94ff91d27 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -217,12 +217,6 @@ def predict_denoiser_output(self, noisy_latent, timestep_or_sigma, condition, ** def prepare_infer_latents(self, height, width, generator=None): raise NotImplementedError - def dmd_latent_shape(self, 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 diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py index e3f56311f..e50e0c9ca 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py @@ -1 +1,5 @@ """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 index 1a3b9c7c7..54b91536a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -25,6 +25,8 @@ ) 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): @@ -253,6 +255,17 @@ def denoiser(self): class GenericDistillationCapability(BoundCapability, DistillationCapability): """Distribution-matching operations shared by image flow models.""" + 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 @@ -293,11 +306,9 @@ def latent_shape( 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}].") - spatial_scale = int(self.model.vae_scale_factor) - size_multiple = spatial_scale * 2 - if height % size_multiple or width % size_multiple: - raise ValueError(f"Image DMD target_height and target_width must be divisible by the model image size multiple {size_multiple}, got {height}x{width}.") - return self.model.dmd_latent_shape(height, width) + 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): @@ -360,7 +371,7 @@ def predict_guided_velocity( if negative_condition is None: return self.predict_velocity(latents, sigma, condition) - if self.model.cfg_on_denoiser_output(): + if self._guidance_in_denoiser_space: denoiser_input = self.model.prepare_denoiser_input(latents) positive = self.model.denoise( denoiser_input, 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/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index 8a1dc5387..df1c9f626 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -5,6 +5,7 @@ from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER @@ -32,7 +33,14 @@ def register_capabilities(self): ) self.capabilities.register( DistillationCapability, - GenericDistillationCapability(self), + GenericDistillationCapability( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="transformer.config.in_channels", + spatial_downsample_multiplier=2, + ), + guidance_in_denoiser_space=True, + ), ) self.capabilities.register( ConsistencyCapability, @@ -164,14 +172,6 @@ def prepare_infer_latents(self, height, width, generator=None): 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, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return (1, 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 diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 04228aaf3..8de0fd5fc 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -5,6 +5,7 @@ from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, DopsdCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability from lightx2v_train.model_zoo.flux2.capability_adapters.flux2_dopsd_capability import ( @@ -35,7 +36,13 @@ def register_capabilities(self): ) self.capabilities.register( DistillationCapability, - GenericDistillationCapability(self), + GenericDistillationCapability( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="transformer.config.in_channels", + spatial_downsample_multiplier=2, + ), + ), ) self.capabilities.register( ConsistencyCapability, @@ -131,16 +138,6 @@ def encode_prompt_condition(self, prompt): def encode_prompt_text(self, prompt): return self.encode_prompt_condition(prompt) - def dmd_latent_shape(self, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return ( - 1, - 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) 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 index 35fd6dd66..617c09634 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -7,6 +7,7 @@ from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.model_zoo.longcat_image.capability_adapters import LongCatImageConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER @@ -34,7 +35,13 @@ def register_capabilities(self): ) self.capabilities.register( DistillationCapability, - GenericDistillationCapability(self), + GenericDistillationCapability( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="vae.config.latent_channels", + ), + guidance_in_denoiser_space=True, + ), ) self.capabilities.register( ConsistencyCapability, @@ -166,14 +173,6 @@ def prepare_infer_latents(self, height, width, generator=None): 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, height, width): - latent_h = 2 * (int(height) // (self.vae_scale_factor * 2)) - latent_w = 2 * (int(width) // (self.vae_scale_factor * 2)) - return (1, 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 diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index ef2acfd6e..9f3f20140 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -5,6 +5,7 @@ from diffusers.image_processor import VaeImageProcessor from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyCapability from lightx2v_train.utils.registry import MODEL_REGISTER @@ -37,7 +38,13 @@ def register_capabilities(self): ) self.capabilities.register( DistillationCapability, - GenericDistillationCapability(self), + GenericDistillationCapability( + self, + latent_geometry=SpatialLatentGeometry( + channels_path="vae.config.z_dim", + temporal_size=1, + ), + ), ) self.capabilities.register( ConsistencyCapability, @@ -155,15 +162,6 @@ def prepare_infer_latents(self, height, width, generator=None): shape = (1, self.vae.config.z_dim, 1, latent_h, latent_w) return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) - def dmd_latent_shape(self, height, width): - return ( - 1, - int(self.vae.config.z_dim), - 1, - int(height) // self.vae_scale_factor, - int(width) // self.vae_scale_factor, - ) - def decode_latent(self, latent): # Reverse the normalization from encode_to_latent: # encode: normalized = (raw - mean) / latents_std diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index 7dead91ce..0cce54550 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -384,13 +384,10 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): def postprocess_denoiser_output(self, prediction, denoiser_input): return prediction - def cfg_on_denoiser_output(self): - return False - def _latent_channels(self): return int(self.transformer.config.in_channels) - def dmd_latent_shape(self, height, width): + def _inference_latent_shape(self, height, width): num_frames = int(self.config.get("inference", {}).get("num_frames", 9)) if num_frames != 1 and (num_frames - 1) % self.vae_scale_factor_temporal != 0: raise ValueError(f"LingBot-Video num_frames must be 1 or 4n+1, got {num_frames}.") @@ -405,7 +402,7 @@ def dmd_latent_shape(self, height, width): ) def prepare_infer_latents(self, height, width, generator=None): - shape = self.dmd_latent_shape(height, width) + shape = self._inference_latent_shape(height, width) return torch.randn(shape, generator=generator, device=self.device, dtype=torch.float32) def _vae_latent_to_dit(self, latents): From 1afa84e119c16043387bcc3bed30bd0b1299f906 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 17 Aug 2026 17:22:26 +0000 Subject: [PATCH 10/34] dopsd --- .../model_capabilities/__init__.py | 10 +- .../model_capabilities/training.py | 128 ++-- .../lightx2v_train/model_zoo/base.py | 58 -- .../model_zoo/capability_adapters/__init__.py | 7 +- .../capability_adapters/peft_adapter_bank.py | 191 +++++ .../flux2/capability_adapters/__init__.py | 3 +- .../flux2_dopsd_capability.py | 284 +++----- .../model_zoo/flux2/flux2_klein.py | 75 +- .../schedulers/flow_matching.py | 22 +- .../lightx2v_train/trainers/dopsd.py | 689 +++++++----------- .../lightx2v_train/trainers/dopsd_core.py | 294 ++++++++ lightx2v_train/tests/test_dopsd_core.py | 193 +++++ .../tests/test_peft_adapter_bank.py | 71 ++ 13 files changed, 1241 insertions(+), 784 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py create mode 100644 lightx2v_train/lightx2v_train/trainers/dopsd_core.py create mode 100644 lightx2v_train/tests/test_dopsd_core.py create mode 100644 lightx2v_train/tests/test_peft_adapter_bank.py diff --git a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py index 1d4b132b1..00ab6f55e 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -6,14 +6,15 @@ ModelCapability, ) from .training import ( + AdapterBankCapability, AutoregressiveDistillationCapability, AutoregressiveRolloutContext, CheckpointCapability, ConsistencyCapability, DistillationCapability, DopsdCapability, - DopsdLossResult, - DopsdStepContext, + DopsdPreparedBatch, + DopsdPreparedTeacherBatch, FlowMatchingSFTCapability, LossResult, ParallelCapability, @@ -25,6 +26,7 @@ ) __all__ = [ + "AdapterBankCapability", "AutoregressiveDistillationCapability", "AutoregressiveRolloutContext", "BoundCapability", @@ -35,8 +37,8 @@ "ConsistencyCapability", "DistillationCapability", "DopsdCapability", - "DopsdLossResult", - "DopsdStepContext", + "DopsdPreparedBatch", + "DopsdPreparedTeacherBatch", "FlowMatchingSFTCapability", "LossResult", "ModelCapability", diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 4a4659804..66d7ea630 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -2,8 +2,9 @@ from abc import abstractmethod from collections.abc import Collection +from contextlib import AbstractContextManager from dataclasses import dataclass, field -from typing import Any, Callable, Mapping +from typing import Any, Callable, Iterable, Mapping import torch from torch import Tensor @@ -232,111 +233,122 @@ def rollout( @dataclass(frozen=True) -class DopsdStepContext: - scheduler: Any - num_training_steps: int - running_dtype: torch.dtype - student_adapter: str - teacher_adapter: str - teacher_prompts: Callable[[Any], list[str]] - step_loss_weights: Callable[[int], list[float]] - collect_trajectory: bool = False +class DopsdPreparedBatch: + """Model-specific values needed by the model-agnostic DOPSD objective.""" + + initial_state: Tensor + state_ids: Tensor + student_condition: Any + teacher_condition: Any + teacher_reference: Any + latent_hw: tuple[int, int] @dataclass(frozen=True) -class DopsdLossResult: - loss: Tensor - student_trajectory: tuple[Tensor, ...] = () - teacher_trajectory: tuple[Tensor, ...] = () - latent_ids: Tensor | None = None - height: int | None = None - width: int | None = None +class DopsdPreparedTeacherBatch: + """Prepared inputs for a privileged-teacher rollout.""" + + initial_state: Tensor + state_ids: Tensor + condition: Any + reference: Any + latent_hw: tuple[int, int] + height: int + width: int class DopsdCapability(ModelCapability): + """Model operations required by DOPSD. + + Adapter lifecycle and objective math intentionally live outside this + capability so a new model only needs to describe its own data path. + """ + @property @abstractmethod - def device(self): + def device(self) -> torch.device: pass @abstractmethod - def configure_adapters( + def prepare_training_batch( self, - rank, - alpha, - target_modules, - student_adapter, - teacher_adapter, - initialize_teacher, - ) -> None: - pass - - @abstractmethod - def parameters(self): + batch: Mapping[str, Any], + teacher_prompts: list[str], + running_dtype: torch.dtype, + ) -> DopsdPreparedBatch: pass @abstractmethod - def compute_loss( + def prepare_teacher_batch( self, - batch, - context: DopsdStepContext, - ) -> DopsdLossResult: + reference_image: Tensor, + teacher_prompts: list[str], + running_dtype: torch.dtype, + generator: torch.Generator | None = None, + ) -> DopsdPreparedTeacherBatch: pass @abstractmethod - def ema_update(self, student_adapter, teacher_adapter, decay) -> None: + def predict_velocity( + self, + state: Tensor, + time: Tensor, + condition: Any, + state_ids: Tensor, + reference: Any = None, + ) -> Tensor: pass @abstractmethod - def decode_trajectory(self, trajectory, latent_ids): + def decode_state(self, state: Tensor, state_ids: Tensor): pass - @abstractmethod - def set_training(self, student_adapter, teacher_adapter) -> None: - pass - @abstractmethod - def set_eval(self) -> None: - pass +class AdapterBankCapability(ModelCapability): + """Lifecycle for multiple named parameter-efficient adapters.""" @abstractmethod - def set_active_adapter(self, adapter_name) -> None: + def configure_pair( + self, + rank: int, + alpha: int, + target_modules: Any, + student_adapter: str, + teacher_adapter: str, + initialize_teacher: bool, + ) -> None: pass @abstractmethod - def encode_prompt(self, prompts): + def parameters(self, adapter_name: str) -> Iterable[Tensor]: pass @abstractmethod - def prepare_reference(self, image): + def activate( + self, + adapter_name: str, + training: bool | None = None, + ) -> AbstractContextManager[None]: pass @abstractmethod - def initial_latents(self, height, width, generator=None): + def set_trainable(self, adapter_name: str) -> None: pass @abstractmethod - def predict_velocity( - self, - latents, - time, - condition, - latent_ids, - adapter_name, - **kwargs, - ): + def copy(self, source_adapter: str, target_adapter: str) -> None: pass @abstractmethod - def load_adapter(self, path, adapter_name, weights_subdir=None) -> None: + def ema_update(self, source_adapter: str, target_adapter: str, decay: float) -> None: pass @abstractmethod - def save_adapter(self, path, adapter_name, weights_subdir=None) -> None: + def load(self, path, adapter_name: str, weights_subdir: str | None = None) -> None: pass @abstractmethod - def copy_adapter(self, source_adapter, target_adapter) -> None: + def save(self, path, adapter_name: str, weights_subdir: str | None = None) -> None: pass diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 94ff91d27..96e5d4245 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -65,64 +65,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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py index e50e0c9ca..42959861a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/__init__.py @@ -1,5 +1,10 @@ """Capability adapters shared by multiple model families.""" from .latent_geometry import LatentGeometry, SpatialLatentGeometry +from .peft_adapter_bank import PeftAdapterBankCapability -__all__ = ["LatentGeometry", "SpatialLatentGeometry"] +__all__ = [ + "LatentGeometry", + "PeftAdapterBankCapability", + "SpatialLatentGeometry", +] diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py new file mode 100644 index 000000000..d403d250c --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +from contextlib import contextmanager + +import torch +from peft import LoraConfig + +from lightx2v_train.model_capabilities import AdapterBankCapability, BoundCapability + + +class PeftAdapterBankCapability(BoundCapability, AdapterBankCapability): + """Strict named-adapter lifecycle for PEFT-compatible denoisers.""" + + def __init__(self, model) -> None: + super().__init__(model) + self._active_adapter: str | None = None + self._known_adapters: set[str] = set() + self._parameter_pair_cache = {} + + @staticmethod + def _validate_name(name: str) -> str: + name = str(name).strip() + if not name or "." in name: + raise ValueError(f"Adapter names must be non-empty and may not contain '.', got {name!r}.") + return name + + @staticmethod + def _adapter_token(adapter_name: str) -> str: + return f".{adapter_name}." + + def _denoiser(self): + return self.model.denoiser_module() + + def _named_adapter_parameters(self, adapter_name: str) -> dict[str, torch.nn.Parameter]: + adapter_name = self._validate_name(adapter_name) + token = self._adapter_token(adapter_name) + parameters = { + name: parameter + for name, parameter in self._denoiser().named_parameters() + if token in f".{name}." + } + if not parameters: + raise RuntimeError(f"No parameters found for adapter {adapter_name!r}.") + self._known_adapters.add(adapter_name) + return parameters + + def _parameter_pairs( + self, + source_adapter: str, + target_adapter: str, + ) -> tuple[tuple[torch.nn.Parameter, torch.nn.Parameter], ...]: + source_adapter = self._validate_name(source_adapter) + target_adapter = self._validate_name(target_adapter) + if source_adapter == target_adapter: + raise ValueError("Source and target adapters must be different.") + cache_key = (source_adapter, target_adapter) + cached = self._parameter_pair_cache.get(cache_key) + if cached is not None: + return cached + + source_token = self._adapter_token(source_adapter) + target_token = self._adapter_token(target_adapter) + all_parameters = dict(self._denoiser().named_parameters()) + source_parameters = self._named_adapter_parameters(source_adapter) + target_parameters = self._named_adapter_parameters(target_adapter) + + pairs = [] + mapped_target_names = set() + for source_name, source_parameter in source_parameters.items(): + padded_name = f".{source_name}." + if padded_name.count(source_token) != 1: + raise RuntimeError(f"Adapter token is ambiguous in parameter name {source_name!r}.") + target_name = padded_name.replace(source_token, target_token, 1)[1:-1] + target_parameter = all_parameters.get(target_name) + if target_parameter is None: + raise RuntimeError(f"Adapter {target_adapter!r} is missing parameter {target_name!r}.") + if source_parameter.shape != target_parameter.shape: + raise RuntimeError( + f"Adapter parameter shape mismatch: {source_name}={tuple(source_parameter.shape)}, " + f"{target_name}={tuple(target_parameter.shape)}." + ) + mapped_target_names.add(target_name) + pairs.append((source_parameter, target_parameter)) + + unexpected_targets = set(target_parameters) - mapped_target_names + if unexpected_targets: + raise RuntimeError( + f"Adapter {target_adapter!r} has parameters without a source match: " + f"{sorted(unexpected_targets)}" + ) + result = tuple(pairs) + self._parameter_pair_cache[cache_key] = result + return result + + def configure_pair( + self, + rank: int, + alpha: int, + target_modules, + student_adapter: str, + teacher_adapter: str, + initialize_teacher: bool, + ) -> None: + student_adapter = self._validate_name(student_adapter) + teacher_adapter = self._validate_name(teacher_adapter) + if student_adapter == teacher_adapter: + raise ValueError("DOPSD student and teacher adapter names must be different.") + if int(rank) <= 0 or int(alpha) <= 0: + raise ValueError(f"LoRA rank and alpha must be positive, got rank={rank}, alpha={alpha}.") + + adapter_config = LoraConfig( + r=int(rank), + lora_alpha=int(alpha), + init_lora_weights="gaussian", + target_modules=target_modules, + ) + denoiser = self._denoiser() + denoiser.requires_grad_(False) + denoiser.add_adapter(adapter_config, adapter_name=student_adapter) + denoiser.add_adapter(adapter_config, adapter_name=teacher_adapter) + self._known_adapters.update((student_adapter, teacher_adapter)) + self._set_active(student_adapter) + if initialize_teacher: + self.copy(student_adapter, teacher_adapter) + self.set_trainable(student_adapter) + # Parallel wrappers may replace parameter objects after configuration. + self._parameter_pair_cache.clear() + + def parameters(self, adapter_name: str): + self._parameter_pair_cache.clear() + return tuple(self._named_adapter_parameters(adapter_name).values()) + + def _set_active(self, adapter_name: str) -> None: + adapter_name = self._validate_name(adapter_name) + if adapter_name not in self._known_adapters: + self._named_adapter_parameters(adapter_name) + self._denoiser().set_adapter(adapter_name) + self._active_adapter = adapter_name + + @contextmanager + def activate(self, adapter_name: str, training: bool | None = None): + denoiser = self._denoiser() + previous_adapter = self._active_adapter + previous_training = denoiser.training + self._set_active(adapter_name) + mode_changed = training is not None and training != previous_training + if mode_changed: + denoiser.train(training) + try: + yield + finally: + if mode_changed: + denoiser.train(previous_training) + if previous_adapter is not None and previous_adapter != adapter_name: + self._set_active(previous_adapter) + + def set_trainable(self, adapter_name: str) -> None: + trainable_parameters = self._named_adapter_parameters(adapter_name) + denoiser = self._denoiser() + denoiser.requires_grad_(False) + denoiser.train() + for parameter in trainable_parameters.values(): + parameter.requires_grad_(True) + self._set_active(adapter_name) + + @torch.no_grad() + def copy(self, source_adapter: str, target_adapter: str) -> None: + for source_parameter, target_parameter in self._parameter_pairs(source_adapter, target_adapter): + target_parameter.copy_(source_parameter) + + @torch.no_grad() + def ema_update(self, source_adapter: str, target_adapter: str, decay: float) -> None: + decay = float(decay) + if not 0.0 <= decay < 1.0: + raise ValueError(f"EMA decay must be in [0, 1), got {decay}.") + for source_parameter, target_parameter in self._parameter_pairs(source_adapter, target_adapter): + target_parameter.mul_(decay).add_(source_parameter, alpha=1.0 - decay) + + def load(self, path, adapter_name: str, weights_subdir: str | None = None) -> None: + self.model.load_lora_weights_for_resume( + path, + adapter_name=self._validate_name(adapter_name), + weights_subdir=weights_subdir, + ) + + def save(self, path, adapter_name: str, weights_subdir: str | None = None) -> None: + self.model.save_lora_weights( + path, + adapter_name=self._validate_name(adapter_name), + weights_subdir=weights_subdir, + ) 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 index 8150f7e40..783f85be4 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -1,5 +1,6 @@ """Capabilities implemented specifically by Flux2 models.""" from .flux2_consistency_capability import Flux2ConsistencyCapability +from .flux2_dopsd_capability import Flux2DopsdCapability -__all__ = ["Flux2ConsistencyCapability"] +__all__ = ["Flux2ConsistencyCapability", "Flux2DopsdCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py index 82f4941b0..a2a62eba9 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py @@ -1,201 +1,131 @@ -"""DOPSD capability for Flux2 models.""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping import torch -import torch.nn.functional as F +from diffusers import Flux2KleinPipeline from lightx2v_train.model_capabilities import ( BoundCapability, DopsdCapability, - DopsdLossResult, - DopsdStepContext, + DopsdPreparedBatch, + DopsdPreparedTeacherBatch, ) +@dataclass(frozen=True) +class Flux2DopsdReference: + latents: torch.Tensor + ids: torch.Tensor + + class Flux2DopsdCapability(BoundCapability, DopsdCapability): + """Flux2 tensor preparation and denoising required by DOPSD.""" + @property - def device(self): + def device(self) -> torch.device: return self.model.device - def configure_adapters( + def _image(self, image: torch.Tensor, running_dtype: torch.dtype) -> torch.Tensor: + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4 or image.shape[0] != 1: + raise ValueError(f"DOPSD requires one reference image per rank with shape [1, C, H, W], got {tuple(image.shape)}.") + return image.to(device=self.device, dtype=running_dtype) + + def _initial_state( self, - rank, - alpha, - target_modules, - student_adapter, - teacher_adapter, - initialize_teacher, - ) -> None: - self.model.add_dual_lora( - rank, - alpha, - target_modules, - student_adapter=student_adapter, - teacher_adapter=teacher_adapter, - init_teacher_from_student=initialize_teacher, + height: int, + width: int, + generator: torch.Generator | None, + ) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int]]: + unpacked_state = self.model.prepare_infer_latents(height, width, generator=generator) + state_ids = Flux2KleinPipeline._prepare_latent_ids(unpacked_state).to(self.device) + state = Flux2KleinPipeline._pack_latents(unpacked_state) + return state, state_ids, tuple(unpacked_state.shape[-2:]) + + def _reference(self, image: torch.Tensor) -> Flux2DopsdReference: + encoded = self.model.vae.encode(image).latent_dist.mode() + normalized = self.model._normalize_patch_latents(encoded) + reference_ids = Flux2KleinPipeline._prepare_image_ids([normalized[:1]]) + reference_ids = reference_ids.repeat(normalized.shape[0], 1, 1).to(self.device) + reference_latents = Flux2KleinPipeline._pack_latents(normalized) + return Flux2DopsdReference( + latents=reference_latents.to(device=self.device, dtype=self.model.running_dtype), + ids=reference_ids, ) - self.set_training(student_adapter, teacher_adapter) - - def parameters(self): - return self.model.trainable_parameters() - def compute_loss( + def prepare_training_batch( self, - batch, - context: DopsdStepContext, - ) -> DopsdLossResult: - image = batch["inputs"].get("target_image") + batch: Mapping[str, Any], + teacher_prompts: list[str], + running_dtype: torch.dtype, + ) -> DopsdPreparedBatch: + image = batch.get("inputs", {}).get("target_image") if image is None: - raise ValueError("D-OPSD requires inputs.target_image.") - image = image.to( - device=self.device, - dtype=context.running_dtype, - ) - if image.shape[0] != 1: - raise ValueError("D-OPSD only supports physical batch size 1.") - _, _, height, width = image.shape - latent_hw = (height // 16, width // 16) - timestep_scale = float(context.scheduler.num_train_timesteps) - - with torch.no_grad(): - student_condition = self.model.encode_condition(batch) - teacher_condition = self.model.encode_prompt_text(context.teacher_prompts(batch["conditioning"]["prompt"])) - reference_latents, reference_ids = self.model.prepare_reference_image_latents(image) - initial_latents, latent_ids = self.model.prepare_dopsd_initial_latents( - height, - width, - ) - context.scheduler.set_timesteps( - context.num_training_steps, - latent_hw=latent_hw, - ) - timesteps = context.scheduler.infer_timesteps - - latents = initial_latents - total_loss = 0.0 - weights = context.step_loss_weights(len(timesteps)) - student_trajectory = [] - teacher_trajectory = [] - for step_index, timestep in enumerate(timesteps): - time = (timestep.reshape(1) / timestep_scale).to( - device=self.device, - dtype=context.running_dtype, - ) - if step_index + 1 < len(timesteps): - next_time = timesteps[step_index + 1].reshape(1) - next_time = next_time / timestep_scale - else: - next_time = torch.zeros_like(time) - next_time = next_time.to( - device=self.device, - dtype=context.running_dtype, - ) - latents = latents.detach().requires_grad_(True) - with torch.no_grad(): - teacher_velocity = self.model.predict_velocity( - latents, - time, - teacher_condition, - latent_ids, - context.teacher_adapter, - teacher_image_latents=reference_latents, - teacher_image_latent_ids=reference_ids, - ) - teacher_x0 = latents - time.reshape(1, 1, 1) * teacher_velocity - student_velocity = self.model.predict_velocity( - latents, - time, - student_condition, - latent_ids, - context.student_adapter, - ) - student_x0 = latents - time.reshape(1, 1, 1) * student_velocity - latents = latents + student_velocity * (next_time - time).reshape(1, 1, 1) - total_loss = total_loss + weights[step_index] * F.mse_loss( - student_x0, - teacher_x0.detach(), - ) - if context.collect_trajectory: - student_trajectory.append(student_x0.detach()) - teacher_trajectory.append(teacher_x0.detach()) - - return DopsdLossResult( - loss=total_loss / sum(weights), - student_trajectory=tuple(student_trajectory), - teacher_trajectory=tuple(teacher_trajectory), - latent_ids=latent_ids if context.collect_trajectory else None, - height=height if context.collect_trajectory else None, - width=width if context.collect_trajectory else None, - ) - - def ema_update(self, student_adapter, teacher_adapter, decay) -> None: - self.model.ema_update_lora_adapter( - student_adapter, - teacher_adapter, - decay, - ) - - def decode_trajectory(self, trajectory, latent_ids): - return self.model.decode_packed_x0_to_images(trajectory, latent_ids) - - def set_training(self, student_adapter, teacher_adapter) -> None: - self.model.set_dual_lora_trainable( - student_adapter, - teacher_adapter, + raise ValueError("DOPSD requires inputs.target_image.") + image = self._image(image, running_dtype) + height, width = image.shape[-2:] + initial_state, state_ids, latent_hw = self._initial_state(height, width, generator=None) + return DopsdPreparedBatch( + initial_state=initial_state, + state_ids=state_ids, + student_condition=self.model.encode_condition(batch), + teacher_condition=self.model.encode_prompt_condition(teacher_prompts), + teacher_reference=self._reference(image), + latent_hw=latent_hw, ) - def set_eval(self) -> None: - self.model.set_denoiser_eval() - - def set_active_adapter(self, adapter_name) -> None: - self.model.set_active_adapter(adapter_name) - - def encode_prompt(self, prompts): - return self.model.encode_prompt_text(prompts) - - def prepare_reference(self, image): - return self.model.prepare_reference_image_latents(image) - - def initial_latents(self, height, width, generator=None): - return self.model.prepare_dopsd_initial_latents( - height, - width, - generator=generator, + def prepare_teacher_batch( + self, + reference_image: torch.Tensor, + teacher_prompts: list[str], + running_dtype: torch.dtype, + generator: torch.Generator | None = None, + ) -> DopsdPreparedTeacherBatch: + image = self._image(reference_image, running_dtype) + height, width = image.shape[-2:] + initial_state, state_ids, latent_hw = self._initial_state(height, width, generator) + return DopsdPreparedTeacherBatch( + initial_state=initial_state, + state_ids=state_ids, + condition=self.model.encode_prompt_condition(teacher_prompts), + reference=self._reference(image), + latent_hw=latent_hw, + height=height, + width=width, ) def predict_velocity( self, - latents, - time, - condition, - latent_ids, - adapter_name, - **kwargs, - ): - return self.model.predict_velocity( - latents, - time, - condition, - latent_ids, - adapter_name, - **kwargs, - ) - - def load_adapter(self, path, adapter_name, weights_subdir=None) -> None: - self.model.load_lora_weights_for_resume( - path, - adapter_name=adapter_name, - weights_subdir=weights_subdir, - ) - - def save_adapter(self, path, adapter_name, weights_subdir=None) -> None: - self.model.save_lora_weights( - path, - adapter_name=adapter_name, - weights_subdir=weights_subdir, - ) - - def copy_adapter(self, source_adapter, target_adapter) -> None: - self.model.copy_lora_adapter_weights( - source_adapter, - target_adapter, - ) + state: torch.Tensor, + time: torch.Tensor, + condition: Any, + state_ids: torch.Tensor, + reference: Any = None, + ) -> torch.Tensor: + hidden_states = state + image_ids = state_ids + if reference is not None: + if not isinstance(reference, Flux2DopsdReference): + raise TypeError(f"Flux2 DOPSD expected Flux2DopsdReference, got {type(reference).__name__}.") + hidden_states = torch.cat([state, reference.latents], dim=1) + image_ids = torch.cat([state_ids, reference.ids], dim=1) + + velocity = self.model.transformer( + hidden_states=hidden_states, + timestep=time, + guidance=None, + encoder_hidden_states=condition["prompt_embed"], + txt_ids=condition["text_ids"], + img_ids=image_ids, + joint_attention_kwargs={}, + return_dict=False, + )[0] + return velocity[:, : state.shape[1]] + + @torch.no_grad() + def decode_state(self, state: torch.Tensor, state_ids: torch.Tensor): + unpacked = Flux2KleinPipeline._unpack_latents_with_ids(state, state_ids) + return self.model.decode_latent(unpacked) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 8de0fd5fc..22ec01ac7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -4,11 +4,20 @@ from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, DopsdCapability, FlowMatchingSFTCapability -from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry +from lightx2v_train.model_capabilities import ( + AdapterBankCapability, + ConsistencyCapability, + DistillationCapability, + DopsdCapability, + FlowMatchingSFTCapability, +) +from lightx2v_train.model_zoo.capability_adapters import ( + PeftAdapterBankCapability, + SpatialLatentGeometry, +) from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability -from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability -from lightx2v_train.model_zoo.flux2.capability_adapters.flux2_dopsd_capability import ( +from lightx2v_train.model_zoo.flux2.capability_adapters import ( + Flux2ConsistencyCapability, Flux2DopsdCapability, ) from lightx2v_train.utils.registry import MODEL_REGISTER @@ -52,6 +61,10 @@ def register_capabilities(self): DopsdCapability, Flux2DopsdCapability(self), ) + self.capabilities.register( + AdapterBankCapability, + PeftAdapterBankCapability(self), + ) def load_components(self, transformer_only=False, reference_model=None): if transformer_only: @@ -135,9 +148,6 @@ def encode_prompt_condition(self, prompt): ) return {"prompt_embed": prompt_embed, "text_ids": text_ids} - def encode_prompt_text(self, prompt): - return self.encode_prompt_condition(prompt) - 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) @@ -180,12 +190,6 @@ def decode_latent(self, 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, @@ -206,48 +210,3 @@ def get_pipeline_infer_kwargs(self, infer_config): "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, generator=None): - pipeline = self.assemble_pipeline() - num_latents_channels = self.transformer.config.in_channels // 4 - return pipeline.prepare_latents( - batch_size=1, - 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/schedulers/flow_matching.py b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py index 958227e2c..e3ee864ca 100644 --- a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py +++ b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py @@ -102,17 +102,29 @@ def _expand_to_ndim(self, values, ndim): # ============================== # The following methods are for inference only # ============================== - def set_timesteps(self, num_inference_steps, sigmas=None, latent_hw=None): - self.num_inference_steps = num_inference_steps - + def build_inference_sigmas(self, num_inference_steps, sigmas=None, latent_hw=None): + """Build an immutable inference schedule without changing scheduler state.""" + num_inference_steps = int(num_inference_steps) + if num_inference_steps <= 0: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}.") if sigmas is None: sigmas = torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) if self.do_time_shift: sigmas = self.time_shift(sigmas, latent_hw=latent_hw, num_steps=num_inference_steps) else: sigmas = torch.tensor(sigmas, dtype=torch.float32) - self.infer_sigmas = torch.cat([sigmas, torch.zeros(1)]).to(self.device) - self.infer_timesteps = (sigmas * self.num_train_timesteps).to(self.device) + if sigmas.ndim != 1 or sigmas.numel() != num_inference_steps: + raise ValueError(f"sigmas must contain exactly {num_inference_steps} values, got shape {tuple(sigmas.shape)}.") + return torch.cat([sigmas, torch.zeros(1)]).to(self.device) + + def set_timesteps(self, num_inference_steps, sigmas=None, latent_hw=None): + self.num_inference_steps = int(num_inference_steps) + self.infer_sigmas = self.build_inference_sigmas( + self.num_inference_steps, + sigmas=sigmas, + latent_hw=latent_hw, + ) + self.infer_timesteps = self.infer_sigmas[:-1] * self.num_train_timesteps def step(self, model_output, step_index, latent): f""" diff --git a/lightx2v_train/lightx2v_train/trainers/dopsd.py b/lightx2v_train/lightx2v_train/trainers/dopsd.py index 8e77ef070..61f85e48b 100644 --- a/lightx2v_train/lightx2v_train/trainers/dopsd.py +++ b/lightx2v_train/lightx2v_train/trainers/dopsd.py @@ -1,28 +1,37 @@ -### -### Source idea: D-OPSD Trainer (https://github.com/vvvvvjdy/D-OPSD) -### +"""Decoupled DOPSD trainer. + +Algorithm math lives in :mod:`dopsd_core`; model families only implement the +DopsdCapability data path, and adapter lifecycle is provided independently by +AdapterBankCapability. +""" + +from __future__ import annotations + import os import shutil from pathlib import Path import torch import torch.distributed.checkpoint as dcp -from diffusers.optimization import get_scheduler from loguru import logger -from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict from lightx2v_train.infer import build_inferencer from lightx2v_train.infer.dopsd_trajectory_viz import save_student_teacher_trajectory_grid -from lightx2v_train.model_capabilities import ( - DopsdCapability, - DopsdStepContext, +from lightx2v_train.model_capabilities import AdapterBankCapability, DopsdCapability +from lightx2v_train.runtime.checkpoint import prune_checkpoints +from lightx2v_train.runtime.distributed import ( + barrier, + get_rank, + get_world_size, + is_distributed, + is_main_process, + reduce_mean, ) -from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints -from lightx2v_train.runtime.distributed import barrier, get_rank, get_world_size, is_distributed, is_main_process, reduce_mean from lightx2v_train.utils.registry import TRAINER_REGISTER -from lightx2v_train.utils.utils import get_running_dtype from .base import BaseTrainer +from .dopsd_core import DopsdConfig, DopsdObjective @TRAINER_REGISTER("dopsd") @@ -30,209 +39,132 @@ class DopsdTrainer(BaseTrainer): required_capabilities = ( *BaseTrainer.required_capabilities, DopsdCapability, + AdapterBankCapability, ) def __init__(self, config): super().__init__(config) - self.running_dtype = get_running_dtype(self.model_config["running_dtype"]) - - lora_config = self.training_config.get("lora", {}) - self.lora_rank = lora_config.get("rank", 16) - self.lora_alpha = lora_config.get("alpha", self.lora_rank) - self.lora_target_modules = lora_config.get("target_modules") - - dopsd_config = self.training_config.get("dopsd", {}) - self.num_training_steps = dopsd_config.get("num_training_steps", 4) - self.ema_decay = dopsd_config.get("ema_decay", 0.999) - self.ema_decay_warmup = dopsd_config.get("ema_decay_warmup") - self.ema_decay_warmup_iters = int(dopsd_config.get("ema_decay_warmup_iters", 0)) - self.step_loss_weights = dopsd_config.get("step_loss_weights") - self.student_adapter = dopsd_config.get("student_adapter", "student") - self.teacher_adapter = dopsd_config.get("teacher_adapter", "teacher") - self.edit_sys_prompt = dopsd_config.get( - "edit_sys_prompt", - "The output must be exactly the same as the reference image.", - ) - self.teacher_use_dataset_prompt = dopsd_config.get("teacher_use_dataset_prompt", False) + self.dopsd_config = DopsdConfig.from_training_config(self.training_config) self._train_reference_map = None - self.trajectory_every_iters = dopsd_config.get("trajectory_every_iters", None) - if self.trajectory_every_iters is not None: - self.trajectory_every_iters = int(self.trajectory_every_iters) - - self.gradient_checkpointing = self.training_config.get("gradient_checkpointing", True) - - optimizer_config = self.training_config.get("optimizer", {}) - self.optimizer_learning_rate = optimizer_config.get("learning_rate", 1e-4) - self.optimizer_adam_beta1 = optimizer_config.get("adam_beta1", 0.9) - self.optimizer_adam_beta2 = optimizer_config.get("adam_beta2", 0.999) - self.optimizer_weight_decay = optimizer_config.get("weight_decay", 0.01) - self.optimizer_adam_epsilon = optimizer_config.get("adam_epsilon", 1e-8) - - self.lr_scheduler_name = self.training_config.get("lr_scheduler", "constant") - self.lr_warmup_iters = self.training_config["lr_warmup_iters"] - self.max_train_iters = self.training_config["max_train_iters"] - - self.output_train_dir = self.training_config["output_dir"] - self.gradient_accumulation_iters = self.training_config["gradient_accumulation_iters"] - self.max_grad_norm = self.training_config.get("max_grad_norm", 1.0) - self.save_every_iters = self.training_config["save_every_iters"] - self.save_total_limit = self.training_config["save_total_limit"] - - self.infer_every_iters = self.infer_config.get("infer_every_iters", None) - logging_config = self.config.get("logging", {}) - self.train_log_every_iters = max(1, int(logging_config.get("train_log_every_iters", 10))) - - resume_config = self.config.get("resume", {}) - self.auto_resume = resume_config.get("auto_resume", False) + if int(self.gradient_accumulation_iters) <= 0: + raise ValueError("training.gradient_accumulation_iters must be positive.") + inference_lora = self.infer_config.get("lora_config") or {} + if inference_lora.get("path"): + raise ValueError("DOPSD inference uses the in-memory student and teacher adapters; inference.lora_config.path must not be set.") def set_model(self, model): super().set_model(model) - self.dopsd = model.capabilities.require(DopsdCapability) - - def _resolve_ema_decay(self, current_iter): - if self.ema_decay_warmup is not None and self.ema_decay_warmup_iters > 0 and current_iter <= self.ema_decay_warmup_iters: - return float(self.ema_decay_warmup) - return float(self.ema_decay) - - def _step_loss_weights_for(self, num_steps): - if not self.step_loss_weights: - return [1.0] * num_steps - weights = [float(w) for w in self.step_loss_weights] - if len(weights) < num_steps: - weights.extend([weights[-1]] * (num_steps - len(weights))) - return weights[:num_steps] + capabilities = model.ensure_capabilities() + self.dopsd = capabilities.require(DopsdCapability) + self.adapters = capabilities.require(AdapterBankCapability) + self.objective = DopsdObjective( + self.dopsd, + self.adapters, + self.noise_scheduler, + self.dopsd_config, + self.running_dtype, + ) def setup(self, resume_ckpt_path=None): - self.dopsd.configure_adapters( - self.lora_rank, - self.lora_alpha, - self.lora_target_modules, - self.student_adapter, - self.teacher_adapter, - resume_ckpt_path is None, + algorithm = self.dopsd_config + self.adapters.configure_pair( + rank=self.lora_rank, + alpha=self.lora_alpha, + target_modules=self.lora_target_modules, + student_adapter=algorithm.student_adapter, + teacher_adapter=algorithm.teacher_adapter, + initialize_teacher=resume_ckpt_path is None, ) - self.parallel.apply(self.config) if self.gradient_checkpointing: self.trainable_model.enable_gradient_checkpointing() - if self.infer_every_iters: self.inferencer = build_inferencer(self.config) self.inferencer.set_model(self.model) self.trainable_model.log_structure() - - self.trainable_params = list(self.dopsd.parameters()) - self.optimizer = torch.optim.AdamW( - self.trainable_params, - lr=self.optimizer_learning_rate, - betas=(self.optimizer_adam_beta1, self.optimizer_adam_beta2), - weight_decay=self.optimizer_weight_decay, - eps=self.optimizer_adam_epsilon, - ) - self.lr_scheduler = get_scheduler( - self.lr_scheduler_name, - optimizer=self.optimizer, - num_warmup_steps=self.lr_warmup_iters, - num_training_steps=self.max_train_iters, - ) + self.trainable_params = list(self.adapters.parameters(algorithm.student_adapter)) + if not self.trainable_params: + raise RuntimeError("DOPSD student adapter has no trainable parameters.") + self.optimizer = self._build_optimizer(self.trainable_params) + self.lr_scheduler = self._build_lr_scheduler(self.optimizer) if resume_ckpt_path is not None: self._load_resume_state(resume_ckpt_path) + self.adapters.set_trainable(algorithm.student_adapter) - def _teacher_lora_checkpoint_path(self, resume_ckpt_path): - return os.path.join(resume_ckpt_path, "teacher", "pytorch_lora_weights.safetensors") + @staticmethod + def _adapter_weights_path(checkpoint_path, weights_subdir=None): + weights_dir = Path(checkpoint_path) + if weights_subdir: + weights_dir /= weights_subdir + return weights_dir / "pytorch_lora_weights.safetensors" def _load_resume_state(self, resume_ckpt_path): if self.parallel.is_fsdp(): - self._load_distributed_state(resume_ckpt_path) + super()._load_distributed_state(resume_ckpt_path) else: self._load_single_process_state(resume_ckpt_path) - teacher_weights_path = self._teacher_lora_checkpoint_path(resume_ckpt_path) - if os.path.exists(teacher_weights_path): - self.dopsd.load_adapter( + algorithm = self.dopsd_config + teacher_path = self._adapter_weights_path(resume_ckpt_path, "teacher") + if teacher_path.exists(): + self.adapters.load( resume_ckpt_path, - self.teacher_adapter, + algorithm.teacher_adapter, weights_subdir="teacher", ) - logger.info("Restored teacher EMA LoRA from {}", teacher_weights_path) + logger.info("Restored teacher EMA adapter from {}", teacher_path) else: - self.dopsd.copy_adapter( - self.student_adapter, - self.teacher_adapter, - ) + self.adapters.copy(algorithm.student_adapter, algorithm.teacher_adapter) logger.warning( - "Teacher LoRA not found in checkpoint {}; initialized teacher from student", + "Teacher adapter was not found in {}; initialized it from the student adapter.", resume_ckpt_path, ) def _load_single_process_state(self, resume_ckpt_path): - training_state_path = os.path.join(resume_ckpt_path, "training_state.pt") - if not os.path.exists(training_state_path): + state_path = Path(resume_ckpt_path) / "training_state.pt" + if not state_path.exists(): raise RuntimeError(f"training_state.pt not found in {resume_ckpt_path}") - - state = torch.load(training_state_path, map_location="cpu", weights_only=False) - self._validate_checkpoint_metadata(state, training_state_path, resume_ckpt_path) - self.dopsd.load_adapter( - resume_ckpt_path, - self.student_adapter, - ) + state = torch.load(state_path, map_location="cpu", weights_only=False) + self._validate_checkpoint_metadata(state, str(state_path), resume_ckpt_path) + self.adapters.load(resume_ckpt_path, self.dopsd_config.student_adapter) self.optimizer.load_state_dict(state["optimizer"]) self.lr_scheduler.load_state_dict(state["lr_scheduler"]) - logger.info("Restored training state from {}", training_state_path) - - def _load_distributed_state(self, resume_ckpt_path): - dist_state_path = os.path.join(resume_ckpt_path, "dist_state") - if not os.path.exists(dist_state_path): - raise RuntimeError(f"FSDP2 resume requires dist_state/, but it was not found in {resume_ckpt_path}") - - trainer_state_path = os.path.join(resume_ckpt_path, "trainer_state.pt") - if not os.path.exists(trainer_state_path): - raise RuntimeError(f"trainer_state.pt not found in {resume_ckpt_path}") - trainer_state = torch.load(trainer_state_path, map_location="cpu", weights_only=False) - self._validate_checkpoint_metadata(trainer_state, trainer_state_path, resume_ckpt_path) - - options = StateDictOptions(ignore_frozen_params=True, strict=False) - state_module = self.parallel.state_module() - model_state, optim_state = get_state_dict(state_module, self.optimizer, options=options) - state = {"model": model_state, "optimizer": optim_state} - dcp.load(state, checkpoint_id=dist_state_path) - set_state_dict( - state_module, - self.optimizer, - model_state_dict=state["model"], - optim_state_dict=state["optimizer"], - options=options, - ) - - self.lr_scheduler.load_state_dict(trainer_state["lr_scheduler"]) - logger.info("Restored distributed training state from {}", resume_ckpt_path) + logger.info("Restored DOPSD training state from {}", state_path) + + def _checkpoint_metadata(self): + metadata = self.dopsd_config.checkpoint_metadata() + target_modules = self.lora_target_modules + if isinstance(target_modules, (list, tuple)): + target_modules = list(target_modules) + metadata["lora"] = { + "rank": int(self.lora_rank), + "alpha": int(self.lora_alpha), + "target_modules": target_modules, + } + return metadata - def _teacher_edit_prompts(self, prompts): - if isinstance(prompts, str): - base_prompts = [prompts] - else: - base_prompts = list(prompts) - suffix = self.edit_sys_prompt.strip() - if not suffix: - return base_prompts - if not self.teacher_use_dataset_prompt: - return [suffix] * len(base_prompts) - return [f"{prompt} {suffix}".strip() for prompt in base_prompts] + def _validate_checkpoint_metadata(self, state, state_path, resume_ckpt_path): + super()._validate_checkpoint_metadata(state, state_path, resume_ckpt_path) + saved_metadata = state.get("dopsd") + if saved_metadata is None: + logger.warning("DOPSD checkpoint {} has no algorithm metadata; treating it as a legacy checkpoint.", state_path) + return + expected_metadata = self._checkpoint_metadata() + if saved_metadata != expected_metadata: + raise RuntimeError(f"DOPSD checkpoint configuration does not match the current run: saved={saved_metadata}, current={expected_metadata}.") def _build_train_reference_map(self): if self._train_reference_map is not None: return self._train_reference_map = {} - train_dataset = self.dataloader_train.dataset - for record in train_dataset.samples: + for record in getattr(self.dataloader_train.dataset, "samples", ()): target_image = record.get("target_image") prompt = record.get("prompt") - if target_image is None or prompt is None: - continue - self._train_reference_map[str(prompt).strip()] = target_image + if target_image is not None and prompt is not None: + self._train_reference_map[str(prompt).strip()] = target_image def _resolve_teacher_reference_image(self, record, dataset): target_image = record.get("target_image") @@ -243,204 +175,146 @@ def _resolve_teacher_reference_image(self, record, dataset): self._build_train_reference_map() fallback_path = self._train_reference_map.get(str(record.get("prompt", "")).strip()) - if fallback_path is None: - return None - return dataset.load_image(fallback_path) - - def _validate_checkpoint_metadata(self, state, state_path, resume_ckpt_path): - checkpoint_world_size = state.get("world_size") - current_world_size = get_world_size() - if checkpoint_world_size != current_world_size: - raise RuntimeError(f"Cannot resume checkpoint saved with world_size={checkpoint_world_size} using world_size={current_world_size}: {state_path}") - - expected_iteration = parse_checkpoint_iteration(resume_ckpt_path) - checkpoint_iteration = state.get("iteration") - if checkpoint_iteration != expected_iteration: - raise RuntimeError(f"Cannot resume checkpoint with iteration={checkpoint_iteration} in {state_path}, expected iteration={expected_iteration} from {resume_ckpt_path}") + return None if fallback_path is None else dataset.load_image(fallback_path) def compute_loss_on_sample(self, sample, collect_trajectory=False): - result = self.dopsd.compute_loss( - sample, - DopsdStepContext( - scheduler=self.noise_scheduler, - num_training_steps=self.num_training_steps, - running_dtype=self.running_dtype, - student_adapter=self.student_adapter, - teacher_adapter=self.teacher_adapter, - teacher_prompts=self._teacher_edit_prompts, - step_loss_weights=self._step_loss_weights_for, - collect_trajectory=collect_trajectory, - ), - ) - if collect_trajectory: - return ( - result.loss, - result.student_trajectory, - result.teacher_trajectory, - result.latent_ids, - result.height, - result.width, - ) - return result.loss + return self.objective.compute_loss(sample, collect_trajectory=collect_trajectory) @torch.no_grad() - def _save_training_trajectory(self, current_iter, student_x0_traj, teacher_x0_traj, latent_ids): - traj_dir = Path(self.output_train_dir) / "trajectory" / f"iter-{current_iter:09d}" - traj_dir.mkdir(parents=True, exist_ok=True) + def _save_training_trajectory(self, current_iter, result): + if result.state_ids is None: + raise RuntimeError("DOPSD trajectory result is missing state IDs.") + trajectory_dir = Path(self.output_train_dir) / "trajectory" / f"iter-{current_iter:09d}" + trajectory_dir.mkdir(parents=True, exist_ok=True) if torch.cuda.is_available(): torch.cuda.empty_cache() + student_images = [] + teacher_images = [] + num_steps = len(result.student_trajectory) + for step_index, (student_x0, teacher_x0) in enumerate(zip(result.student_trajectory, result.teacher_trajectory)): + logger.info( + "[train] trajectory decode iter={} step={}/{}", + current_iter, + step_index + 1, + num_steps, + ) + student_images.extend(self.dopsd.decode_state(student_x0, result.state_ids)) + teacher_images.extend(self.dopsd.decode_state(teacher_x0, result.state_ids)) - student_step_images = [] - teacher_step_images = [] - num_steps = len(student_x0_traj) - for step_idx, (x_0_student, x_0_teacher) in enumerate(zip(student_x0_traj, teacher_x0_traj)): - logger.info("[train] trajectory decode iter={} step={}/{} student", current_iter, step_idx + 1, num_steps) - student_step_images.extend(self.dopsd.decode_trajectory(x_0_student, latent_ids)) - logger.info("[train] trajectory decode iter={} step={}/{} teacher", current_iter, step_idx + 1, num_steps) - teacher_step_images.extend(self.dopsd.decode_trajectory(x_0_teacher, latent_ids)) - - save_path = traj_dir / "student_teacher_x0_traj.png" - save_student_teacher_trajectory_grid(student_step_images, teacher_step_images, save_path) + save_path = trajectory_dir / "student_teacher_x0_traj.png" + save_student_teacher_trajectory_grid(student_images, teacher_images, save_path) logger.info("[train] saved trajectory iter={} path={}", current_iter, save_path) - def _resolve_resume(self): - if not self.auto_resume: - return None, 0 - ckpt_path, current_iter = find_latest_checkpoint(self.output_train_dir) - if ckpt_path is None: - logger.info("Auto-resume enabled but no checkpoint found in '{}'. Starting from scratch.", self.output_train_dir) - else: - logger.info("Auto-resuming from checkpoint: {} (iteration {})", ckpt_path, current_iter) - return ckpt_path, current_iter - def train(self): - resume_ckpt_path, current_iter = self._resolve_resume() - self.setup(resume_ckpt_path=resume_ckpt_path) + resume_path, current_iter = self._resolve_resume() + self.setup(resume_ckpt_path=resume_path) if is_main_process(): os.makedirs(self.output_train_dir, exist_ok=True) barrier() - max_train_iters = self.max_train_iters - grad_accum_iters = self.gradient_accumulation_iters - max_grad_norm = self.max_grad_norm - save_every_iters = self.save_every_iters - save_total_limit = self.save_total_limit + algorithm = self.dopsd_config grad_accum_counter = 0 running_loss = 0.0 - logger.info( - "[train] dopsd start iter={}/{} world_size={} grad_accum={} num_training_steps={} " - "ema_decay={} ema_decay_warmup={} ema_decay_warmup_iters={} step_loss_weights={} " - "edit_sys_prompt={!r} teacher_use_dataset_prompt={} trajectory_every_iters={}", + "[train] dopsd start iter={}/{} world_size={} grad_accum={} steps={} " + "ema={} warmup_ema={} warmup_iters={} weights={} teacher_prompt={!r} " + "teacher_use_dataset_prompt={} trajectory_every_iters={}", current_iter, - max_train_iters, + self.max_train_iters, get_world_size(), - grad_accum_iters, - self.num_training_steps, - self.ema_decay, - self.ema_decay_warmup, - self.ema_decay_warmup_iters, - self._step_loss_weights_for(self.num_training_steps), - self.edit_sys_prompt, - self.teacher_use_dataset_prompt, - self.trajectory_every_iters, + self.gradient_accumulation_iters, + algorithm.num_training_steps, + algorithm.ema_decay, + algorithm.ema_decay_warmup, + algorithm.ema_decay_warmup_iters, + algorithm.loss_weights(algorithm.num_training_steps), + algorithm.edit_sys_prompt, + algorithm.teacher_use_dataset_prompt, + algorithm.trajectory_every_iters, ) + if self.infer_every_iters: self.inferencer.set_data(self.dataloader_eval) if current_iter == 0: self.run_inference(current_iter) epoch = 0 - while current_iter < max_train_iters: + while current_iter < self.max_train_iters: sampler = getattr(self.dataloader_train, "sampler", None) if hasattr(sampler, "set_epoch"): sampler.set_epoch(epoch) for sample in self.dataloader_train: - sync_grad = (grad_accum_counter + 1) % grad_accum_iters == 0 + sync_grad = (grad_accum_counter + 1) % self.gradient_accumulation_iters == 0 self._set_gradient_sync(sync_grad) - - should_save_trajectory = self.trajectory_every_iters and (current_iter + 1) % self.trajectory_every_iters == 0 - if should_save_trajectory: - loss, student_x0_traj, teacher_x0_traj, latent_ids, _height, _width = self.compute_loss_on_sample( - sample, - collect_trajectory=True, - ) - else: - loss = self.compute_loss_on_sample(sample) - (loss / grad_accum_iters).backward() - running_loss += loss.item() / grad_accum_iters + collect_trajectory = bool(sync_grad and algorithm.trajectory_every_iters and (current_iter + 1) % algorithm.trajectory_every_iters == 0) + result = self.compute_loss_on_sample(sample, collect_trajectory=collect_trajectory) + (result.loss / self.gradient_accumulation_iters).backward() + running_loss += result.loss.item() / self.gradient_accumulation_iters grad_accum_counter += 1 - if grad_accum_counter % grad_accum_iters != 0: + if not sync_grad: continue - torch.nn.utils.clip_grad_norm_(self.trainable_params, max_grad_norm) + torch.nn.utils.clip_grad_norm_(self.trainable_params, self.max_grad_norm) self.optimizer.step() self.lr_scheduler.step() self.optimizer.zero_grad() - current_ema_decay = self._resolve_ema_decay(current_iter + 1) - self.dopsd.ema_update( - self.student_adapter, - self.teacher_adapter, - current_ema_decay, + ema_decay = algorithm.ema_decay_at(current_iter + 1) + self.adapters.ema_update( + algorithm.student_adapter, + algorithm.teacher_adapter, + ema_decay, ) - current_iter += 1 + display_loss = reduce_mean(running_loss) current_lr = self.lr_scheduler.get_last_lr()[0] - if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= max_train_iters: + if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= self.max_train_iters: logger.info( "[train] iter={}/{} loss_dopsd={:.6f} lr={:.8f} ema_decay={:.4f}", current_iter, - max_train_iters, + self.max_train_iters, display_loss, current_lr, - current_ema_decay, + ema_decay, ) self.log_metrics( { "train/loss_dopsd": display_loss, "train/lr": current_lr, - "train/ema_decay": current_ema_decay, + "train/ema_decay": ema_decay, }, step=current_iter, ) running_loss = 0.0 - if should_save_trajectory: + if collect_trajectory: + self.trainable_model.set_eval() barrier() if is_main_process(): - logger.info("[train] saving trajectory iter={} (decoding {} x0 pairs)...", current_iter, len(student_x0_traj)) - self.dopsd.set_eval() - self._save_training_trajectory( - current_iter, - student_x0_traj, - teacher_x0_traj, - latent_ids, - ) - self.dopsd.set_training( - self.student_adapter, - self.teacher_adapter, - ) + self._save_training_trajectory(current_iter, result) barrier() + self.adapters.set_trainable(algorithm.student_adapter) - if save_every_iters and current_iter % save_every_iters == 0: - self.save_checkpoint(current_iter, save_total_limit) - + if self.save_every_iters and current_iter % self.save_every_iters == 0: + self.save_checkpoint(current_iter, self.save_total_limit) if self.infer_every_iters and current_iter % self.infer_every_iters == 0: self.run_inference(current_iter) - - if current_iter >= max_train_iters: + if current_iter >= self.max_train_iters: break - epoch += 1 - logger.info("[train] finished iter={}/{}", current_iter, max_train_iters) + self.finish_monitor() + logger.info("[train] finished iter={}/{}", current_iter, self.max_train_iters) - def _set_gradient_sync(self, enabled): - self.parallel.set_gradient_sync(enabled) + def _first_teacher_reference(self, dataset, samples): + for record in samples: + reference = self._resolve_teacher_reference_image(record, dataset) + if reference is not None: + return reference, record + return None, None @torch.no_grad() def _run_teacher_inference(self, current_iter, iter_output_dir): @@ -448,17 +322,20 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): samples = dataset.samples rank = get_rank() world_size = get_world_size() + num_steps = int(self.infer_config.get("num_inference_steps", self.dopsd_config.num_training_steps)) + if num_steps <= 0: + raise ValueError("inference.num_inference_steps must be positive.") + base_seed = int(self.infer_config.get("seed", 42)) + + teacher_output_dir = Path(iter_output_dir) / "teacher" + teacher_output_dir.mkdir(parents=True, exist_ok=True) + dummy_reference, dummy_record = self._first_teacher_reference(dataset, samples) + if dummy_reference is None: + logger.warning("[train] skipped teacher inference: evaluation data has no reference images.") + barrier() + return - num_inference_steps = self.infer_config.get("num_inference_steps", self.num_training_steps) - base_seed = self.infer_config.get("seed", 42) - t_scale = float(self.noise_scheduler.num_train_timesteps) - - teacher_output_dir = os.path.join(iter_output_dir, "teacher") - os.makedirs(teacher_output_dir, exist_ok=True) - - self.dopsd.set_active_adapter(self.teacher_adapter) - self.dopsd.set_eval() - + self.trainable_model.set_eval() num_slots = (len(samples) + world_size - 1) // world_size if is_distributed() else len(samples) saved_count = 0 skipped_count = 0 @@ -466,102 +343,65 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): "[train] running teacher inference iter={} output_dir={} steps={}", current_iter, teacher_output_dir, - num_inference_steps, + num_steps, ) - for slot in range(num_slots): - i = slot * world_size + rank if is_distributed() else slot - if i >= len(samples): - continue - record = samples[i] - reference_image = self._resolve_teacher_reference_image(record, dataset) - if reference_image is None: + for slot in range(num_slots): + sample_index = slot * world_size + rank if is_distributed() else slot + has_sample = sample_index < len(samples) + record = samples[sample_index] if has_sample else dummy_record + reference = self._resolve_teacher_reference_image(record, dataset) if has_sample else dummy_reference + should_save = has_sample and reference is not None + if reference is None: + reference = dummy_reference + if has_sample and not should_save: skipped_count += 1 logger.warning( - "[train] teacher infer skip sample={}/{}: missing target_image in eval metadata", - i + 1, + "[train] teacher infer skip sample={}/{}: missing target_image in evaluation metadata", + sample_index + 1, len(samples), ) - continue - - image = reference_image.unsqueeze(0).to( - device=self.dopsd.device, - dtype=self.running_dtype, - ) - height, width = image.shape[2], image.shape[3] - latent_hw = (height // 16, width // 16) - - teacher_prompts = self._teacher_edit_prompts(record["prompt"]) - teacher_condition = self.dopsd.encode_prompt(teacher_prompts) - teacher_image_latents, teacher_image_latent_ids = self.dopsd.prepare_reference(image) - seed = base_seed + i + if not torch.is_tensor(reference): + raise TypeError(f"DOPSD teacher references must be tensors after dataset loading, got {type(reference).__name__}.") + reference = reference.unsqueeze(0) if reference.ndim == 3 else reference + prompt = record.get("prompt", "") + teacher_prompts = self.dopsd_config.teacher_prompts(prompt) + seed = base_seed + sample_index if has_sample else base_seed generator = torch.Generator(device=self.dopsd.device).manual_seed(seed) - latents, latent_ids = self.dopsd.initial_latents( - height, - width, + rollout = self.objective.rollout_teacher( + reference, + teacher_prompts, + num_steps, generator=generator, ) - - self.noise_scheduler.set_timesteps(num_inference_steps, latent_hw=latent_hw) - timesteps = self.noise_scheduler.infer_timesteps - num_steps = len(timesteps) + if not should_save: + continue logger.info( "[train] teacher infer sample={}/{} seed={} size={}x{} prompt={!r}", - i + 1, + sample_index + 1, len(samples), seed, - height, - width, + rollout.height, + rollout.width, teacher_prompts[0], ) - for back_step in range(num_steps): - t = timesteps[back_step].expand(1) / t_scale - t = t.to(device=self.dopsd.device, dtype=self.running_dtype) - if back_step < num_steps - 1: - next_t = timesteps[back_step + 1].expand(1) / t_scale - else: - next_t = torch.zeros_like(t) - next_t = next_t.to( - device=self.dopsd.device, - dtype=self.running_dtype, - ) - dt = next_t - t - - v_pred = self.dopsd.predict_velocity( - latents, - t, - teacher_condition, - latent_ids, - self.teacher_adapter, - teacher_image_latents=teacher_image_latents, - teacher_image_latent_ids=teacher_image_latent_ids, - ) - latents = latents + v_pred * dt.reshape(1, 1, 1) - - images = self.dopsd.decode_trajectory(latents, latent_ids) - save_path = Path(teacher_output_dir) / f"{i:05d}.png" + images = self.dopsd.decode_state(rollout.state, rollout.state_ids) + save_path = teacher_output_dir / f"{sample_index:05d}.png" images[0].save(save_path) saved_count += 1 - logger.info("[train] teacher infer sample={}/{} saved path={}", i + 1, len(samples), save_path) + logger.info("[train] teacher infer sample={}/{} saved path={}", sample_index + 1, len(samples), save_path) barrier() if is_distributed(): - saved_tensor = torch.tensor( - saved_count, - device=self.dopsd.device, - dtype=torch.int64, - ) - skipped_tensor = torch.tensor( - skipped_count, + counts = torch.tensor( + [saved_count, skipped_count], device=self.dopsd.device, dtype=torch.int64, ) - torch.distributed.all_reduce(saved_tensor, op=torch.distributed.ReduceOp.SUM) - torch.distributed.all_reduce(skipped_tensor, op=torch.distributed.ReduceOp.SUM) - saved_count = saved_tensor.item() - skipped_count = skipped_tensor.item() + torch.distributed.all_reduce(counts, op=torch.distributed.ReduceOp.SUM) + saved_count, skipped_count = counts.tolist() logger.info( "[train] finished teacher inference iter={} saved={} skipped={}", current_iter, @@ -570,84 +410,89 @@ def _run_teacher_inference(self, current_iter, iter_output_dir): ) def run_inference(self, current_iter): - base_output_dir = self.infer_config.get("output_dir", "./output_infer") - iter_output_dir = os.path.join(base_output_dir, f"iter-{current_iter:09d}") - - self.dopsd.set_active_adapter(self.student_adapter) - self.inferencer.output_infer_dir = iter_output_dir - os.makedirs(iter_output_dir, exist_ok=True) - logger.info("[train] running student inference iter={} output_dir={}", current_iter, iter_output_dir) - self.inferencer.infer() - barrier() - self._run_teacher_inference(current_iter, iter_output_dir) + output_dir = Path(self.infer_config.get("output_dir", "./output_infer")) / f"iter-{current_iter:09d}" + output_dir.mkdir(parents=True, exist_ok=True) + self.inferencer.output_infer_dir = str(output_dir) + logger.info("[train] running student inference iter={} output_dir={}", current_iter, output_dir) + with self.adapters.activate(self.dopsd_config.student_adapter): + self.inferencer.infer() + self._run_teacher_inference(current_iter, output_dir) + self.adapters.set_trainable(self.dopsd_config.student_adapter) logger.info("[train] finished inference iter={}", current_iter) - self.dopsd.set_training( - self.student_adapter, - self.teacher_adapter, - ) + def _training_state(self, iteration, include_optimizer): + state = { + "iteration": iteration, + "world_size": get_world_size(), + "lr_scheduler": self.lr_scheduler.state_dict(), + "dopsd": self._checkpoint_metadata(), + } + if include_optimizer: + state["optimizer"] = self.optimizer.state_dict() + return state + + @staticmethod + def _atomic_torch_save(state, path): + path = Path(path) + temporary_path = path.with_suffix(path.suffix + ".tmp") + torch.save(state, temporary_path) + os.replace(temporary_path, path) def save_checkpoint(self, iteration, save_total_limit): if is_main_process(): prune_checkpoints(self.output_train_dir, save_total_limit) - - save_dir = os.path.join(self.output_train_dir, f"checkpoint-{iteration:09d}") - logger.info("[train] saving checkpoint iter={} path={}", iteration, save_dir) + save_dir = Path(self.output_train_dir) / f"checkpoint-{iteration:09d}" + logger.info("[train] saving DOPSD checkpoint iter={} path={}", iteration, save_dir) if is_main_process(): - os.makedirs(save_dir, exist_ok=True) + save_dir.mkdir(parents=True, exist_ok=True) barrier() - self.dopsd.save_adapter(save_dir, self.student_adapter) + self.adapters.save(save_dir, self.dopsd_config.student_adapter) barrier() - self.dopsd.save_adapter( - save_dir, - self.teacher_adapter, - weights_subdir="teacher", - ) + self.adapters.save(save_dir, self.dopsd_config.teacher_adapter, weights_subdir="teacher") barrier() config_path = self.config.get("config_path") if is_main_process() and config_path is not None: - shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) + shutil.copy2(config_path, save_dir / "config.yaml") if self.parallel.is_fsdp(): self._save_distributed_state(save_dir, iteration) + else: + if is_main_process(): + self._atomic_torch_save( + self._training_state(iteration, include_optimizer=True), + save_dir / "training_state.pt", + ) barrier() - logger.info("[train] saved checkpoint iter={} path={} (student + teacher EMA LoRA)", iteration, save_dir) - return - training_state = { - "iteration": iteration, - "world_size": get_world_size(), - "optimizer": self.optimizer.state_dict(), - "lr_scheduler": self.lr_scheduler.state_dict(), - } if is_main_process(): - torch.save(training_state, os.path.join(save_dir, "training_state.pt")) + (save_dir / "_SUCCESS").touch() barrier() - logger.info("[train] saved checkpoint iter={} path={} (student + teacher EMA LoRA)", iteration, save_dir) + logger.info("[train] saved DOPSD checkpoint iter={} path={}", iteration, save_dir) def _save_distributed_state(self, save_dir, iteration): - dist_state_path = os.path.join(save_dir, "dist_state") + dist_state_path = Path(save_dir) / "dist_state" if is_main_process(): - os.makedirs(dist_state_path, exist_ok=True) - torch.save( - { - "iteration": iteration, - "world_size": get_world_size(), - "lr_scheduler": self.lr_scheduler.state_dict(), - }, - os.path.join(save_dir, "trainer_state.pt"), - ) + dist_state_path.mkdir(parents=True, exist_ok=True) barrier() options = StateDictOptions(ignore_frozen_params=True, strict=False) - model_state, optim_state = get_state_dict( + model_state, optimizer_state = get_state_dict( self.parallel.state_module(), self.optimizer, options=options, ) dcp.save( - {"model": model_state, "optimizer": optim_state}, - checkpoint_id=dist_state_path, + {"model": model_state, "optimizer": optimizer_state}, + checkpoint_id=str(dist_state_path), ) + barrier() + + # trainer_state.pt is a runtime completion marker, so write it last. + if is_main_process(): + self._atomic_torch_save( + self._training_state(iteration, include_optimizer=False), + Path(save_dir) / "trainer_state.pt", + ) + barrier() diff --git a/lightx2v_train/lightx2v_train/trainers/dopsd_core.py b/lightx2v_train/lightx2v_train/trainers/dopsd_core.py new file mode 100644 index 000000000..384e250e6 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/dopsd_core.py @@ -0,0 +1,294 @@ +"""DOPSD configuration and objective, independent of any model family. + +The implementation follows the privileged editing-teacher / text-only +student setup introduced by https://github.com/vvvvvjdy/D-OPSD. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Mapping + +import torch +from torch import Tensor + +from lightx2v_train.model_capabilities import ( + AdapterBankCapability, + DopsdCapability, +) + + +@dataclass(frozen=True) +class DopsdConfig: + num_training_steps: int = 4 + ema_decay: float = 0.999 + ema_decay_warmup: float | None = None + ema_decay_warmup_iters: int = 0 + step_loss_weights: tuple[float, ...] | None = None + student_adapter: str = "student" + teacher_adapter: str = "teacher" + edit_sys_prompt: str = "The output must be exactly the same as the reference image." + teacher_use_dataset_prompt: bool = False + trajectory_every_iters: int | None = None + + @classmethod + def from_training_config(cls, training_config: Mapping[str, Any]) -> "DopsdConfig": + train_type = training_config.get("train_type") + if train_type != "lora": + raise ValueError( + "DOPSD trains a student LoRA and an EMA teacher LoRA; " + f"training.train_type must be 'lora', got {train_type!r}." + ) + raw = training_config.get("dopsd", {}) or {} + if not isinstance(raw, Mapping): + raise ValueError("training.dopsd must be a mapping.") + + raw_weights = raw.get("step_loss_weights") + if raw_weights is None: + weights = None + else: + if isinstance(raw_weights, (str, bytes, Mapping)): + raise ValueError("training.dopsd.step_loss_weights must be a sequence of numbers.") + try: + weights = tuple(float(value) for value in raw_weights) + except TypeError as error: + raise ValueError("training.dopsd.step_loss_weights must be a sequence of numbers.") from error + warmup = raw.get("ema_decay_warmup") + trajectory_interval = raw.get("trajectory_every_iters") + config = cls( + num_training_steps=int(raw.get("num_training_steps", 4)), + ema_decay=float(raw.get("ema_decay", 0.999)), + ema_decay_warmup=None if warmup is None else float(warmup), + ema_decay_warmup_iters=int(raw.get("ema_decay_warmup_iters", 0)), + step_loss_weights=weights, + student_adapter=str(raw.get("student_adapter", "student")).strip(), + teacher_adapter=str(raw.get("teacher_adapter", "teacher")).strip(), + edit_sys_prompt=str( + raw.get( + "edit_sys_prompt", + "The output must be exactly the same as the reference image.", + ) + ), + teacher_use_dataset_prompt=bool(raw.get("teacher_use_dataset_prompt", False)), + trajectory_every_iters=None if trajectory_interval is None else int(trajectory_interval), + ) + config.validate() + return config + + def validate(self) -> None: + if self.num_training_steps <= 0: + raise ValueError(f"training.dopsd.num_training_steps must be positive, got {self.num_training_steps}.") + self._validate_decay("ema_decay", self.ema_decay) + if self.ema_decay_warmup is not None: + self._validate_decay("ema_decay_warmup", self.ema_decay_warmup) + if self.ema_decay_warmup_iters < 0: + raise ValueError("training.dopsd.ema_decay_warmup_iters may not be negative.") + if not self.student_adapter or not self.teacher_adapter: + raise ValueError("DOPSD adapter names must be non-empty.") + if self.student_adapter == self.teacher_adapter: + raise ValueError("DOPSD student_adapter and teacher_adapter must be different.") + if "." in self.student_adapter or "." in self.teacher_adapter: + raise ValueError("DOPSD adapter names may not contain '.'.") + if self.trajectory_every_iters is not None and self.trajectory_every_iters <= 0: + raise ValueError("training.dopsd.trajectory_every_iters must be positive when set.") + + if self.step_loss_weights is not None: + if not self.step_loss_weights: + raise ValueError("training.dopsd.step_loss_weights may not be empty.") + if any(not math.isfinite(weight) or weight < 0.0 for weight in self.step_loss_weights): + raise ValueError("training.dopsd.step_loss_weights must be finite and non-negative.") + if sum(self.loss_weights(self.num_training_steps)) <= 0.0: + raise ValueError("DOPSD loss weights must contain at least one positive value.") + + @staticmethod + def _validate_decay(name: str, decay: float) -> None: + if not math.isfinite(decay) or not 0.0 <= decay < 1.0: + raise ValueError(f"training.dopsd.{name} must be in [0, 1), got {decay}.") + + def loss_weights(self, num_steps: int) -> tuple[float, ...]: + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}.") + if self.step_loss_weights is None: + return (1.0,) * num_steps + weights = list(self.step_loss_weights[:num_steps]) + if len(weights) < num_steps: + weights.extend([weights[-1]] * (num_steps - len(weights))) + return tuple(weights) + + def ema_decay_at(self, iteration: int) -> float: + if self.ema_decay_warmup is not None and 0 < iteration <= self.ema_decay_warmup_iters: + return self.ema_decay_warmup + return self.ema_decay + + def teacher_prompts(self, prompts) -> list[str]: + base_prompts = [prompts] if isinstance(prompts, str) else list(prompts) + if not base_prompts: + raise ValueError("DOPSD requires at least one prompt.") + suffix = self.edit_sys_prompt.strip() + if not suffix: + return [str(prompt) for prompt in base_prompts] + if not self.teacher_use_dataset_prompt: + return [suffix] * len(base_prompts) + return [f"{prompt} {suffix}".strip() for prompt in base_prompts] + + def checkpoint_metadata(self) -> dict[str, Any]: + return { + "version": 2, + "num_training_steps": self.num_training_steps, + "student_adapter": self.student_adapter, + "teacher_adapter": self.teacher_adapter, + } + + +@dataclass(frozen=True) +class DopsdRolloutResult: + state: Tensor + state_ids: Tensor + height: int + width: int + + +@dataclass(frozen=True) +class DopsdLossResult: + loss: Tensor + student_trajectory: tuple[Tensor, ...] = () + teacher_trajectory: tuple[Tensor, ...] = () + state_ids: Tensor | None = None + + +class DopsdObjective: + """The model-independent on-policy student/EMA-teacher objective.""" + + def __init__( + self, + model: DopsdCapability, + adapters: AdapterBankCapability, + scheduler, + config: DopsdConfig, + running_dtype: torch.dtype, + ) -> None: + self.model = model + self.adapters = adapters + self.scheduler = scheduler + self.config = config + self.running_dtype = running_dtype + + @staticmethod + def _expand_time(time: Tensor, state: Tensor) -> Tensor: + if time.ndim == 0: + time = time.reshape(1) + if time.shape[0] not in (1, state.shape[0]): + raise ValueError( + f"Time batch dimension must be 1 or {state.shape[0]}, got {tuple(time.shape)}." + ) + return time.reshape(time.shape[0], *([1] * (state.ndim - 1))) + + def _schedule(self, num_steps: int, latent_hw: tuple[int, int], state: Tensor) -> Tensor: + return self.scheduler.build_inference_sigmas( + num_steps, + latent_hw=latent_hw, + ).to(device=state.device, dtype=self.running_dtype) + + def compute_loss(self, batch: Mapping[str, Any], collect_trajectory: bool = False) -> DopsdLossResult: + prompts = batch.get("conditioning", {}).get("prompt") + if prompts is None: + raise ValueError("DOPSD requires conditioning.prompt.") + + with torch.no_grad(): + prepared = self.model.prepare_training_batch( + batch, + self.config.teacher_prompts(prompts), + self.running_dtype, + ) + sigmas = self._schedule( + self.config.num_training_steps, + prepared.latent_hw, + prepared.initial_state, + ) + + state = prepared.initial_state + weights = self.config.loss_weights(sigmas.numel() - 1) + total_loss: Tensor | None = None + student_trajectory = [] + teacher_trajectory = [] + + for step_index, (time, next_time) in enumerate(zip(sigmas[:-1], sigmas[1:])): + time = time.reshape(1) + next_time = next_time.reshape(1) + # Truncate the rollout graph between steps while retaining the + # differentiable student regression performed at every step. + state = state.detach().requires_grad_(True) + + with torch.no_grad(), self.adapters.activate( + self.config.teacher_adapter, + training=False, + ): + teacher_velocity = self.model.predict_velocity( + state, + time, + prepared.teacher_condition, + prepared.state_ids, + reference=prepared.teacher_reference, + ) + teacher_x0 = state - self._expand_time(time, state) * teacher_velocity + + with self.adapters.activate(self.config.student_adapter, training=True): + student_velocity = self.model.predict_velocity( + state, + time, + prepared.student_condition, + prepared.state_ids, + ) + student_x0 = state - self._expand_time(time, state) * student_velocity + step_loss = (student_x0.float() - teacher_x0.float()).square().mean() + weighted_loss = weights[step_index] * step_loss + total_loss = weighted_loss if total_loss is None else total_loss + weighted_loss + state = state + self._expand_time(next_time - time, state) * student_velocity + + if collect_trajectory: + student_trajectory.append(student_x0.detach()) + teacher_trajectory.append(teacher_x0.detach()) + + if total_loss is None: + raise RuntimeError("DOPSD produced an empty rollout.") + return DopsdLossResult( + loss=total_loss / sum(weights), + student_trajectory=tuple(student_trajectory), + teacher_trajectory=tuple(teacher_trajectory), + state_ids=prepared.state_ids if collect_trajectory else None, + ) + + @torch.no_grad() + def rollout_teacher( + self, + reference_image: Tensor, + teacher_prompts: list[str], + num_steps: int, + generator: torch.Generator | None = None, + ) -> DopsdRolloutResult: + prepared = self.model.prepare_teacher_batch( + reference_image, + teacher_prompts, + self.running_dtype, + generator=generator, + ) + sigmas = self._schedule(num_steps, prepared.latent_hw, prepared.initial_state) + state = prepared.initial_state + with self.adapters.activate(self.config.teacher_adapter, training=False): + for time, next_time in zip(sigmas[:-1], sigmas[1:]): + time = time.reshape(1) + velocity = self.model.predict_velocity( + state, + time, + prepared.condition, + prepared.state_ids, + reference=prepared.reference, + ) + state = state + self._expand_time(next_time.reshape(1) - time, state) * velocity + return DopsdRolloutResult( + state=state, + state_ids=prepared.state_ids, + height=prepared.height, + width=prepared.width, + ) diff --git a/lightx2v_train/tests/test_dopsd_core.py b/lightx2v_train/tests/test_dopsd_core.py new file mode 100644 index 000000000..d1c3a29e3 --- /dev/null +++ b/lightx2v_train/tests/test_dopsd_core.py @@ -0,0 +1,193 @@ +from contextlib import contextmanager + +import pytest +import torch +from lightx2v_train.model_capabilities import ( + AdapterBankCapability, + DopsdCapability, + DopsdPreparedBatch, + DopsdPreparedTeacherBatch, +) +from lightx2v_train.schedulers.flow_matching import RectifiedFlowMatchingScheduler +from lightx2v_train.trainers.dopsd_core import DopsdConfig, DopsdObjective + + +class FakeAdapterBank(AdapterBankCapability): + def __init__(self): + self.active = "student" + self.training = True + self.student = torch.nn.Parameter(torch.tensor(1.0)) + self.teacher = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False) + + def configure_pair(self, *args, **kwargs): + pass + + def parameters(self, adapter_name): + return (self.student,) + + @contextmanager + def activate(self, adapter_name, training=None): + previous_adapter = self.active + previous_training = self.training + self.active = adapter_name + if training is not None: + self.training = training + try: + yield + finally: + self.active = previous_adapter + self.training = previous_training + + def set_trainable(self, adapter_name): + self.active = adapter_name + + def copy(self, source_adapter, target_adapter): + self.teacher.copy_(self.student.detach()) + + def ema_update(self, source_adapter, target_adapter, decay): + self.teacher.mul_(decay).add_(self.student.detach(), alpha=1.0 - decay) + + def load(self, path, adapter_name, weights_subdir=None): + pass + + def save(self, path, adapter_name, weights_subdir=None): + pass + + +class FakeDopsdModel(DopsdCapability): + def __init__(self, adapters): + self.adapters = adapters + + @property + def device(self): + return torch.device("cpu") + + @staticmethod + def _initial_state(): + return torch.ones(1, 2, 3, 4) + + def prepare_training_batch(self, batch, teacher_prompts, running_dtype): + return DopsdPreparedBatch( + initial_state=self._initial_state(), + state_ids=torch.zeros(1, 6, 1), + student_condition="student-condition", + teacher_condition="teacher-condition", + teacher_reference="reference", + latent_hw=(3, 4), + ) + + def prepare_teacher_batch(self, reference_image, teacher_prompts, running_dtype, generator=None): + return DopsdPreparedTeacherBatch( + initial_state=self._initial_state(), + state_ids=torch.zeros(1, 6, 1), + condition="teacher-condition", + reference="reference", + latent_hw=(3, 4), + height=24, + width=32, + ) + + def predict_velocity(self, state, time, condition, state_ids, reference=None): + weight = self.adapters.student if self.adapters.active == "student" else self.adapters.teacher + return torch.ones_like(state) * weight + + def decode_state(self, state, state_ids): + return state + + +class FakeScheduler: + def build_inference_sigmas(self, num_steps, latent_hw=None): + assert num_steps == 2 + assert latent_hw == (3, 4) + return torch.tensor([1.0, 0.5, 0.0]) + + +def _config(**overrides): + values = { + "num_training_steps": 2, + "step_loss_weights": (2.0, 1.0), + } + values.update(overrides) + config = DopsdConfig(**values) + config.validate() + return config + + +def test_dopsd_objective_matches_weighted_x0_regression_and_backpropagates(): + adapters = FakeAdapterBank() + objective = DopsdObjective( + FakeDopsdModel(adapters), + adapters, + FakeScheduler(), + _config(), + torch.float32, + ) + result = objective.compute_loss( + {"conditioning": {"prompt": "a prompt"}}, + collect_trajectory=True, + ) + + # teacher velocity=0.25 and student velocity=1.0, so the x0 delta is + # sigma * 0.75 at sigmas 1.0 and 0.5. + assert result.loss.item() == pytest.approx((2.0 * 0.75**2 + 0.375**2) / 3.0) + assert len(result.student_trajectory) == 2 + assert result.state_ids is not None + result.loss.backward() + assert adapters.student.grad.item() == pytest.approx(1.125) + assert adapters.teacher.grad is None + assert adapters.active == "student" + assert adapters.training is True + + +def test_teacher_rollout_uses_the_same_schedule_and_restores_adapter_state(): + adapters = FakeAdapterBank() + objective = DopsdObjective( + FakeDopsdModel(adapters), + adapters, + FakeScheduler(), + _config(), + torch.float32, + ) + rollout = objective.rollout_teacher( + torch.zeros(1, 3, 24, 32), + ["teacher prompt"], + num_steps=2, + ) + + assert torch.allclose(rollout.state, torch.full_like(rollout.state, 0.75)) + assert adapters.active == "student" + assert adapters.training is True + + +@pytest.mark.parametrize( + ("dopsd", "message"), + [ + ({"num_training_steps": 0}, "num_training_steps"), + ({"ema_decay": 1.0}, "ema_decay"), + ({"student_adapter": "same", "teacher_adapter": "same"}, "must be different"), + ({"step_loss_weights": [0.0, 0.0]}, "at least one positive"), + ({"trajectory_every_iters": 0}, "trajectory_every_iters"), + ], +) +def test_dopsd_config_rejects_invalid_values(dopsd, message): + with pytest.raises(ValueError, match=message): + DopsdConfig.from_training_config({"train_type": "lora", "dopsd": dopsd}) + + +def test_dopsd_config_rejects_full_training(): + with pytest.raises(ValueError, match="must be 'lora'"): + DopsdConfig.from_training_config({"train_type": "full"}) + + +def test_scheduler_can_build_a_schedule_without_mutating_inference_state(): + scheduler = RectifiedFlowMatchingScheduler( + { + "model": {"running_dtype": "fp32"}, + "scheduler": {"num_train_timesteps": 1000}, + } + ) + sigmas = scheduler.build_inference_sigmas(2, latent_hw=(3, 4)) + + assert torch.allclose(sigmas.cpu(), torch.tensor([1.0, 0.5, 0.0])) + assert scheduler.infer_sigmas is None + assert scheduler.infer_timesteps is None diff --git a/lightx2v_train/tests/test_peft_adapter_bank.py b/lightx2v_train/tests/test_peft_adapter_bank.py new file mode 100644 index 000000000..a6c51538c --- /dev/null +++ b/lightx2v_train/tests/test_peft_adapter_bank.py @@ -0,0 +1,71 @@ +import torch +from lightx2v_train.model_zoo.capability_adapters import PeftAdapterBankCapability + + +class FakeDenoiser(torch.nn.Module): + def __init__(self): + super().__init__() + self.lora_A = torch.nn.ModuleDict( + { + "student": torch.nn.Linear(2, 1, bias=False), + "teacher": torch.nn.Linear(2, 1, bias=False), + } + ) + self.active_adapter = None + + def set_adapter(self, adapter_name): + self.active_adapter = adapter_name + + +class FakeModel: + def __init__(self): + self.denoiser = FakeDenoiser() + + def denoiser_module(self): + return self.denoiser + + def load_lora_weights_for_resume(self, *args, **kwargs): + pass + + def save_lora_weights(self, *args, **kwargs): + pass + + +def test_adapter_copy_ema_trainability_and_scoped_activation_are_strict(): + model = FakeModel() + adapters = PeftAdapterBankCapability(model) + student = model.denoiser.lora_A["student"].weight + teacher = model.denoiser.lora_A["teacher"].weight + student.data.fill_(2.0) + teacher.data.zero_() + + adapters.copy("student", "teacher") + assert torch.allclose(teacher, torch.full_like(teacher, 2.0)) + + student.data.fill_(4.0) + adapters.ema_update("student", "teacher", decay=0.5) + assert torch.allclose(teacher, torch.full_like(teacher, 3.0)) + + adapters.set_trainable("student") + assert student.requires_grad + assert not teacher.requires_grad + assert model.denoiser.active_adapter == "student" + + with adapters.activate("teacher", training=False): + assert model.denoiser.active_adapter == "teacher" + assert not model.denoiser.training + assert model.denoiser.active_adapter == "student" + assert model.denoiser.training + + +def test_adapter_pairing_fails_instead_of_silently_skipping_missing_weights(): + model = FakeModel() + del model.denoiser.lora_A["teacher"] + adapters = PeftAdapterBankCapability(model) + + try: + adapters.copy("student", "teacher") + except RuntimeError as error: + assert "No parameters found" in str(error) + else: + raise AssertionError("Expected strict adapter pairing to reject a missing teacher.") From 537730e84427974071c22cde928524606f89bf7f Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 17 Aug 2026 17:25:26 +0000 Subject: [PATCH 11/34] ci --- .../capability_adapters/peft_adapter_bank.py | 16 +++------------- .../lightx2v_train/trainers/dopsd_core.py | 18 ++++++++---------- lightx2v_train/tests/test_dopsd_core.py | 1 + lightx2v_train/tests/test_peft_adapter_bank.py | 1 + 4 files changed, 13 insertions(+), 23 deletions(-) diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py index d403d250c..df0c8bc37 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/peft_adapter_bank.py @@ -34,11 +34,7 @@ def _denoiser(self): def _named_adapter_parameters(self, adapter_name: str) -> dict[str, torch.nn.Parameter]: adapter_name = self._validate_name(adapter_name) token = self._adapter_token(adapter_name) - parameters = { - name: parameter - for name, parameter in self._denoiser().named_parameters() - if token in f".{name}." - } + parameters = {name: parameter for name, parameter in self._denoiser().named_parameters() if token in f".{name}."} if not parameters: raise RuntimeError(f"No parameters found for adapter {adapter_name!r}.") self._known_adapters.add(adapter_name) @@ -75,19 +71,13 @@ def _parameter_pairs( if target_parameter is None: raise RuntimeError(f"Adapter {target_adapter!r} is missing parameter {target_name!r}.") if source_parameter.shape != target_parameter.shape: - raise RuntimeError( - f"Adapter parameter shape mismatch: {source_name}={tuple(source_parameter.shape)}, " - f"{target_name}={tuple(target_parameter.shape)}." - ) + raise RuntimeError(f"Adapter parameter shape mismatch: {source_name}={tuple(source_parameter.shape)}, {target_name}={tuple(target_parameter.shape)}.") mapped_target_names.add(target_name) pairs.append((source_parameter, target_parameter)) unexpected_targets = set(target_parameters) - mapped_target_names if unexpected_targets: - raise RuntimeError( - f"Adapter {target_adapter!r} has parameters without a source match: " - f"{sorted(unexpected_targets)}" - ) + raise RuntimeError(f"Adapter {target_adapter!r} has parameters without a source match: {sorted(unexpected_targets)}") result = tuple(pairs) self._parameter_pair_cache[cache_key] = result return result diff --git a/lightx2v_train/lightx2v_train/trainers/dopsd_core.py b/lightx2v_train/lightx2v_train/trainers/dopsd_core.py index 384e250e6..14351b7a4 100644 --- a/lightx2v_train/lightx2v_train/trainers/dopsd_core.py +++ b/lightx2v_train/lightx2v_train/trainers/dopsd_core.py @@ -36,10 +36,7 @@ class DopsdConfig: def from_training_config(cls, training_config: Mapping[str, Any]) -> "DopsdConfig": train_type = training_config.get("train_type") if train_type != "lora": - raise ValueError( - "DOPSD trains a student LoRA and an EMA teacher LoRA; " - f"training.train_type must be 'lora', got {train_type!r}." - ) + raise ValueError(f"DOPSD trains a student LoRA and an EMA teacher LoRA; training.train_type must be 'lora', got {train_type!r}.") raw = training_config.get("dopsd", {}) or {} if not isinstance(raw, Mapping): raise ValueError("training.dopsd must be a mapping.") @@ -179,9 +176,7 @@ def _expand_time(time: Tensor, state: Tensor) -> Tensor: if time.ndim == 0: time = time.reshape(1) if time.shape[0] not in (1, state.shape[0]): - raise ValueError( - f"Time batch dimension must be 1 or {state.shape[0]}, got {tuple(time.shape)}." - ) + raise ValueError(f"Time batch dimension must be 1 or {state.shape[0]}, got {tuple(time.shape)}.") return time.reshape(time.shape[0], *([1] * (state.ndim - 1))) def _schedule(self, num_steps: int, latent_hw: tuple[int, int], state: Tensor) -> Tensor: @@ -220,9 +215,12 @@ def compute_loss(self, batch: Mapping[str, Any], collect_trajectory: bool = Fals # differentiable student regression performed at every step. state = state.detach().requires_grad_(True) - with torch.no_grad(), self.adapters.activate( - self.config.teacher_adapter, - training=False, + with ( + torch.no_grad(), + self.adapters.activate( + self.config.teacher_adapter, + training=False, + ), ): teacher_velocity = self.model.predict_velocity( state, diff --git a/lightx2v_train/tests/test_dopsd_core.py b/lightx2v_train/tests/test_dopsd_core.py index d1c3a29e3..dad31a53e 100644 --- a/lightx2v_train/tests/test_dopsd_core.py +++ b/lightx2v_train/tests/test_dopsd_core.py @@ -2,6 +2,7 @@ import pytest import torch + from lightx2v_train.model_capabilities import ( AdapterBankCapability, DopsdCapability, diff --git a/lightx2v_train/tests/test_peft_adapter_bank.py b/lightx2v_train/tests/test_peft_adapter_bank.py index a6c51538c..1064f94cf 100644 --- a/lightx2v_train/tests/test_peft_adapter_bank.py +++ b/lightx2v_train/tests/test_peft_adapter_bank.py @@ -1,4 +1,5 @@ import torch + from lightx2v_train.model_zoo.capability_adapters import PeftAdapterBankCapability From 0fb2c115e59471e6440a96c1809b0a1b9d75b764 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 18 Aug 2026 11:09:26 +0000 Subject: [PATCH 12/34] update h3 --- .../train/dmd/minimax_h3_t2av_dmd_lora.yaml | 22 +-- .../model_capabilities/__init__.py | 2 + .../model_capabilities/training.py | 17 +++ .../minimax_h3_distillation_capability.py | 63 +++++++-- .../model_zoo/minimax_h3/minimax_h3_t2av.py | 9 +- .../trainers/dmd/autoregressive_dmd.py | 2 +- .../lightx2v_train/trainers/dmd/config.py | 8 ++ .../trainers/dmd/minimax_h3_trainer.py | 79 ----------- .../lightx2v_train/trainers/dmd/runtime.py | 28 +++- .../trainers/dmd/score_sampling.py | 109 +++++++++++++++ .../lightx2v_train/trainers/dmd/trainer.py | 44 +++--- .../trainers/phased_dmd/rollout.py | 30 ++--- .../lightx2v_train/trainers/sgmd.py | 2 +- .../lightx2v_train/utils/registry.py | 1 - .../tests/test_dmd_score_sampling.py | 126 ++++++++++++++++++ 15 files changed, 401 insertions(+), 141 deletions(-) delete mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py create mode 100644 lightx2v_train/lightx2v_train/trainers/dmd/score_sampling.py create mode 100644 lightx2v_train/tests/test_dmd_score_sampling.py 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 170ec9171..a5bfdc26e 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 @@ -14,6 +14,13 @@ model: use_autocast: false # Full SDPA is impractical at H3 sequence lengths; this matches run_h3.py. attention_backend: _flash_3_hub + capabilities: + distillation: + # 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 @@ -51,7 +58,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 @@ -60,13 +67,6 @@ training: lr_warmup_iters: 0 save_every_iters: 100 save_total_limit: 5 - minimax_h3: - # These shifts are applied independently inside the H3 capability. - 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 @@ -75,8 +75,10 @@ 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: diff --git a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py index 00ab6f55e..5122678ae 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -12,6 +12,7 @@ CheckpointCapability, ConsistencyCapability, DistillationCapability, + DistillationProfile, DopsdCapability, DopsdPreparedBatch, DopsdPreparedTeacherBatch, @@ -36,6 +37,7 @@ "CheckpointCapability", "ConsistencyCapability", "DistillationCapability", + "DistillationProfile", "DopsdCapability", "DopsdPreparedBatch", "DopsdPreparedTeacherBatch", diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 66d7ea630..395c82b9d 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -96,9 +96,26 @@ def denoiser(self): """Return the module owned by this capability.""" +@dataclass(frozen=True) +class DistillationProfile: + """Feature and precision constraints published by a distillation adapter.""" + + supports_guidance: bool = True + supports_cdm: 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 + + class DistillationCapability(ModelCapability): """Model operations required by distribution-matching algorithms.""" + @property + def profile(self) -> DistillationProfile: + return DistillationProfile() + @property @abstractmethod def device(self) -> torch.device: diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py index 0f6d63d1b..d09ef3352 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py @@ -2,9 +2,13 @@ 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 DistillationProfile from lightx2v_train.model_zoo.capability_adapters.common import ( GenericDistillationCapability, _require_single_prompt, @@ -20,6 +24,34 @@ from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape +@dataclass(frozen=True) +class MiniMaxH3DistillationOptions: + video_loss_weight: float = 1.0 + audio_loss_weight: float = 1.0 + video_flow_shift: float = 12.0 + audio_flow_shift: float = 3.0 + + @classmethod + def from_mapping(cls, config: Mapping | None) -> "MiniMaxH3DistillationOptions": + if config is None: + config = {} + if not isinstance(config, Mapping): + raise ValueError("model.capabilities.distillation 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", 12.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.""" @@ -49,22 +81,29 @@ class MiniMaxH3DistillationCapability(GenericDistillationCapability): "ff.net.0.proj", "ff.net.2", ) + _PROFILE = DistillationProfile( + supports_guidance=False, + supports_cdm=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) -> None: + def __init__(self, model, options: Mapping | None = None) -> None: super().__init__(model) - config = model.config["training"].get("minimax_h3", {}) - self.video_weight = float(config.get("video_loss_weight", 1.0)) - self.audio_weight = float(config.get("audio_loss_weight", 1.0)) - self.video_shift = float(config.get("video_flow_shift", 12.0)) - self.audio_shift = float(config.get("audio_flow_shift", 3.0)) - if self.video_shift <= 0 or self.audio_shift <= 0: - raise ValueError("MiniMax-H3 video_flow_shift and audio_flow_shift must be positive.") - if self.video_weight < 0 or self.audio_weight < 0: - raise ValueError("MiniMax-H3 video_loss_weight and audio_loss_weight cannot be negative.") - if self.video_weight == 0 and self.audio_weight == 0: - raise ValueError("At least one MiniMax-H3 modality loss weight must be non-zero.") + options = MiniMaxH3DistillationOptions.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) -> DistillationProfile: + return self._PROFILE + @property def default_negative_prompt(self): return "" 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 index 2e9465869..466601230 100644 --- 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 @@ -1,5 +1,6 @@ """LightX2V-Train wrapper for the trainable MiniMax-H3 T2AV DiT.""" +from collections.abc import Mapping from contextlib import nullcontext import torch @@ -22,9 +23,15 @@ class MiniMaxH3T2AVModel(BaseModel): 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.") self.capabilities.register( DistillationCapability, - MiniMaxH3DistillationCapability(self), + MiniMaxH3DistillationCapability( + self, + capability_config.get("distillation"), + ), ) self.capabilities.register( ConsistencyCapability, diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py index 582a841f5..df9dbda48 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py @@ -56,7 +56,7 @@ def run_back_simulation( denoising_steps=self.denoising_steps, denoising_scheduler=self.denoising_scheduler, trajectory_scheduler=self.scheduler, - running_dtype=self.running_dtype, + running_dtype=self.latent_dtype, frames_per_chunk=self.num_frame_per_chunk, same_step_across_blocks=self.same_step_across_blocks, context_noise=self.context_noise, diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/config.py b/lightx2v_train/lightx2v_train/trainers/dmd/config.py index 53162fa86..7436f183b 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/config.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/config.py @@ -6,7 +6,10 @@ Optional, ) +import torch + from lightx2v_train.utils.image_size_buckets import parse_image_size_buckets +from lightx2v_train.utils.utils import get_running_dtype def _lora_config( @@ -46,6 +49,7 @@ class DmdConfig: random_schedule_sigma_min: float random_schedule_sigma_max: float random_schedule_sampling_method: str + latent_dtype: torch.dtype | None @classmethod def from_mapping( @@ -72,6 +76,9 @@ def from_mapping( negative_prompt = default_negative_prompt if default_negative_prompt is not None else configured_negative_prompt random_schedule = dmd.get("random_schedule", {}) + latent_dtype = dmd.get("latent_dtype") + if latent_dtype is not None: + latent_dtype = get_running_dtype(str(latent_dtype).lower()) image_sizes = dmd.get("image_sizes", []) parse_image_size_buckets(image_sizes) return cls( @@ -114,6 +121,7 @@ def from_mapping( "sampling_method", "stratified", ), + latent_dtype=latent_dtype, ) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py deleted file mode 100644 index 7f5cd9d1e..000000000 --- a/lightx2v_train/lightx2v_train/trainers/dmd/minimax_h3_trainer.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Thin DMD trainer specialization for MiniMax-H3 T2AV.""" - -import torch -from loguru import logger - -from lightx2v_train.runtime.sequence_parallel import broadcast_sequence_parallel_value -from lightx2v_train.utils.registry import TRAINER_REGISTER - -from .trainer import DmdTrainer - - -@TRAINER_REGISTER("minimax_h3_t2av_dmd") -class MiniMaxH3T2AVDmdTrainer(DmdTrainer): - """Configure generic DMD for H3's supported training feature set. - - Audio/video tensor math and the clean-ward velocity convention live in - ``MiniMaxH3DistillationCapability``. This class only validates algorithm - options and supplies H3's continuous, unshifted score-sigma sampling. - """ - - trainer_name = "minimax_h3_t2av_dmd" - default_negative_prompt = "" - supports_cdm = False - supports_ida = False - supports_diversity_loss = False - supports_real_data_fake = False - - def __init__(self, config): - super().__init__(config) - model_name = self.model_config.get("name") - if model_name != "minimax_h3_t2av": - raise ValueError(f"minimax_h3_t2av_dmd requires model.name=minimax_h3_t2av, got {model_name!r}.") - if self.guidance_scale != 1.0: - raise ValueError("MiniMax-H3 is guidance-distilled and has no unconditional branch; set training.teacher.guidance_scale to 1.0.") - if self.warp_denoising_step: - raise ValueError("MiniMax-H3 applies separate video/audio flow shifts in its model capability; set training.dmd.warp_denoising_step to false.") - - h3_config = self.training_config.get("minimax_h3", {}) - if not isinstance(h3_config, dict): - raise TypeError("training.minimax_h3 must be a mapping.") - dtype_name = str(h3_config.get("latent_dtype", "fp32")).lower() - dtype_aliases = { - "fp32": torch.float32, - "float32": torch.float32, - "bf16": torch.bfloat16, - "bfloat16": torch.bfloat16, - } - if dtype_name not in dtype_aliases: - supported = ", ".join(sorted(dtype_aliases)) - raise ValueError(f"Unsupported training.minimax_h3.latent_dtype={dtype_name!r}; expected one of: {supported}.") - self.latent_dtype = dtype_aliases[dtype_name] - # Generic DMD uses running_dtype for rollout states and score sigmas. - # H3 defaults to fp32 latents even when prompt embeddings are bf16. - self.running_dtype = self.latent_dtype - - self.renoise_sigma_min = float(self.dmd_config.get("renoise_sigma_min", 0.02)) - self.renoise_sigma_max = float(self.dmd_config.get("renoise_sigma_max", 0.98)) - if not 0.0 <= self.renoise_sigma_min < self.renoise_sigma_max <= 1.0: - raise ValueError(f"MiniMax-H3 renoise sigma range must satisfy 0 <= min < max <= 1, got [{self.renoise_sigma_min}, {self.renoise_sigma_max}].") - logger.info( - "[train] MiniMax-H3 generic DMD latent_dtype={} renoise_sigma=[{}, {}]", - dtype_name, - self.renoise_sigma_min, - self.renoise_sigma_max, - ) - - def _sample_score_sigma( - self, - denoised_timestep_from, - denoised_timestep_to, - device, - dtype, - ): - del denoised_timestep_from, denoised_timestep_to - sigma = torch.empty((1,), device=device, dtype=torch.float32).uniform_( - self.renoise_sigma_min, - self.renoise_sigma_max, - ) - return broadcast_sequence_parallel_value(sigma).to(dtype=dtype) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py index b849fb42e..9a56e618a 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py @@ -99,10 +99,16 @@ def __init__(self, config): self.random_schedule_sigma_min = parsed.random_schedule_sigma_min self.random_schedule_sigma_max = parsed.random_schedule_sigma_max self.random_schedule_sampling_method = parsed.random_schedule_sampling_method + self._configured_latent_dtype = parsed.latent_dtype + self.latent_dtype = parsed.latent_dtype or self.running_dtype def set_model(self, model): super().set_model(model) self.student = model.capabilities.require(DistillationCapability) + profile = self.student.profile + if self._configured_latent_dtype is None and profile.default_latent_dtype is not None: + self.latent_dtype = profile.default_latent_dtype + self._validate_distillation_profile(profile) if self.negative_prompt is None: self.negative_prompt = self.student.default_negative_prompt for lora_config in ( @@ -112,6 +118,24 @@ def set_model(self, model): if lora_config is not None and "target_modules" not in lora_config and self.student.default_lora_target_modules is not None: lora_config["target_modules"] = list(self.student.default_lora_target_modules) + logger.info("[train] dmd latent_dtype={}", self.latent_dtype) + + def _validate_distillation_profile(self, profile): + model_name = self.model_config.get("name", type(self.model).__name__) + requested_features = ( + (self.cdm_trick.enabled, profile.supports_cdm, "CDM"), + (self.ida_trick.enabled, profile.supports_ida, "IDA"), + (getattr(getattr(self, "diversity_trick", None), "enabled", False), profile.supports_diversity, "diversity loss"), + (getattr(getattr(self, "real_data_fake_trick", None), "enabled", False), profile.supports_real_data_fake, "real-data fake loss"), + ) + unsupported = [name for enabled, supported, name in requested_features if enabled and not supported] + if unsupported: + raise ValueError(f"model={model_name!r} does not support DMD features: {', '.join(unsupported)}.") + if not profile.supports_guidance and self.guidance_scale != 1.0: + raise ValueError(f"model={model_name!r} has no unconditional branch; training.teacher.guidance_scale must be 1.0.") + if getattr(self, "warp_denoising_step", False) and not profile.supports_warped_denoising_schedule: + raise ValueError(f"model={model_name!r} requires an unwarped base denoising schedule; set training.dmd.warp_denoising_step=false.") + @property def cdm_enabled(self): return self.cdm_trick.enabled @@ -353,7 +377,7 @@ def _predict_teacher_velocity(self, latents, sigma, condition, negative_conditio def sample_initial_latents(self, latent_shape): return self.student.initial_latents( latent_shape, - self.running_dtype, + self.latent_dtype, broadcast_sequence_parallel_value, ) @@ -378,7 +402,7 @@ def _cdm_step_context( condition=condition, current_iteration=current_iter, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, scheduler=self.scheduler, predict_student_velocity=lambda latent, sigma, cond: self._predict_velocity( self.student, diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/score_sampling.py b/lightx2v_train/lightx2v_train/trainers/dmd/score_sampling.py new file mode 100644 index 000000000..521d7330a --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/dmd/score_sampling.py @@ -0,0 +1,109 @@ +"""Score-noise sampling policies shared by distribution-matching trainers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class ScoreSigmaContext: + denoised_timestep_from: int | None + denoised_timestep_to: int | None + num_train_timesteps: int + device: torch.device + + +class ScoreSigmaSampler(ABC): + @abstractmethod + def sample(self, context: ScoreSigmaContext) -> torch.Tensor: + """Sample one base sigma in noise-ward coordinates.""" + + +@dataclass(frozen=True) +class DiscreteTimestepScoreSigmaSampler(ScoreSigmaSampler): + """Sample a discrete training timestep, preserving the legacy DMD policy.""" + + sample_min_timestep: int + clamp_min_timestep: int + clamp_max_timestep: int + timestep_shift: float = 1.0 + use_rollout_min: bool = False + use_rollout_max: bool = False + + def __post_init__(self) -> None: + if self.sample_min_timestep < 0: + raise ValueError(f"Minimum score sampling timestep cannot be negative, got {self.sample_min_timestep}.") + if self.clamp_min_timestep < 0 or self.clamp_max_timestep <= self.clamp_min_timestep: + raise ValueError(f"Score timestep clamp bounds must satisfy 0 <= min < max, got [{self.clamp_min_timestep}, {self.clamp_max_timestep}].") + if self.timestep_shift <= 0: + raise ValueError(f"Score timestep shift must be positive, got {self.timestep_shift}.") + + def sample(self, context: ScoreSigmaContext) -> torch.Tensor: + lower = context.denoised_timestep_to if self.use_rollout_min and context.denoised_timestep_to is not None else self.sample_min_timestep + upper = context.denoised_timestep_from if self.use_rollout_max and context.denoised_timestep_from is not None else context.num_train_timesteps + lower = max(0, int(lower)) + upper = min(context.num_train_timesteps, int(upper)) + if upper <= lower: + upper = min(context.num_train_timesteps, lower + 1) + if upper <= lower: + raise ValueError(f"No score timestep remains in [{lower}, {upper}) for num_train_timesteps={context.num_train_timesteps}.") + + timestep = torch.randint(lower, upper, (1,), device=context.device, dtype=torch.long).float() + if self.timestep_shift > 1.0: + normalized = timestep / context.num_train_timesteps + timestep = self.timestep_shift * normalized / (1.0 + (self.timestep_shift - 1.0) * normalized) * context.num_train_timesteps + return timestep.clamp(self.clamp_min_timestep, self.clamp_max_timestep) / context.num_train_timesteps + + +@dataclass(frozen=True) +class ContinuousUniformScoreSigmaSampler(ScoreSigmaSampler): + """Sample an unshifted continuous sigma from a uniform distribution.""" + + min_sigma: float = 0.02 + max_sigma: float = 0.98 + + def __post_init__(self) -> None: + if not 0.0 <= self.min_sigma < self.max_sigma <= 1.0: + raise ValueError(f"Score sigma bounds must satisfy 0 <= min < max <= 1, got [{self.min_sigma}, {self.max_sigma}].") + + def sample(self, context: ScoreSigmaContext) -> torch.Tensor: + return torch.empty((1,), device=context.device, dtype=torch.float32).uniform_(self.min_sigma, self.max_sigma) + + +def build_score_sigma_sampler( + config, + *, + sample_min_timestep: int, + clamp_min_timestep: int, + clamp_max_timestep: int, + timestep_shift: float, + use_rollout_min: bool, + use_rollout_max: bool, +) -> ScoreSigmaSampler: + """Build a policy, defaulting to the pre-existing discrete DMD behavior.""" + + if config is None: + config = {} + if not isinstance(config, Mapping): + raise ValueError("training.dmd.score_sampling must be a mapping.") + + kind = str(config.get("type", "discrete_timestep")).lower() + if kind == "discrete_timestep": + return DiscreteTimestepScoreSigmaSampler( + sample_min_timestep=int(config.get("sample_min_timestep", sample_min_timestep)), + clamp_min_timestep=int(config.get("clamp_min_timestep", clamp_min_timestep)), + clamp_max_timestep=int(config.get("clamp_max_timestep", clamp_max_timestep)), + timestep_shift=float(config.get("timestep_shift", timestep_shift)), + use_rollout_min=bool(config.get("use_rollout_min", use_rollout_min)), + use_rollout_max=bool(config.get("use_rollout_max", use_rollout_max)), + ) + if kind == "continuous_uniform": + return ContinuousUniformScoreSigmaSampler( + min_sigma=float(config.get("min_sigma", 0.02)), + max_sigma=float(config.get("max_sigma", 0.98)), + ) + raise ValueError(f"Unsupported training.dmd.score_sampling.type={kind!r}; expected 'discrete_timestep' or 'continuous_uniform'.") diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py index 69ff60b7a..326432f72 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py @@ -39,6 +39,7 @@ from .config import DmdScheduleConfig from .runtime import _DmdRuntime +from .score_sampling import ScoreSigmaContext, build_score_sigma_sampler @TRAINER_REGISTER("dmd") @@ -71,6 +72,15 @@ def __init__(self, config): self.min_score_timestep = parsed.min_score_timestep self.student_checkpoint_path = parsed.student_checkpoint_path self.student_checkpoint_strict = parsed.student_checkpoint_strict + self.score_sigma_sampler = build_score_sigma_sampler( + self.dmd_config.get("score_sampling"), + sample_min_timestep=self.min_score_timestep, + clamp_min_timestep=self.min_step, + clamp_max_timestep=self.max_step, + timestep_shift=self.score_timestep_shift, + use_rollout_min=self.ts_schedule, + use_rollout_max=self.ts_schedule_max, + ) self.diversity_trick = DiversityTrick.from_mapping(self.dmd_config.get("div_loss", {})) self.diversity_trick.validate_trainer( @@ -489,7 +499,7 @@ def _build_denoising_steps(self, device): def _extract_real_latents(self, sample): return self.student.extract_real_latents( sample, - self.running_dtype, + self.latent_dtype, broadcast_sequence_parallel_value, ) @@ -536,7 +546,7 @@ def _diversity_step_context(self, initial_noise, conditions): negative_condition=conditions[1], latent_hw=initial_noise.shape[-2:], device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, predict_teacher_velocity=self._predict_teacher_velocity, predict_student_velocity=partial( self._predict_velocity, @@ -583,7 +593,7 @@ def _real_data_fake_context(self, sample, conditions, region="main"): condition=conditions[0], negative_condition=conditions[1], device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, extract_real_latents=self._extract_real_latents, sample_synced_int=self._sample_synced_int, broadcast_noise=broadcast_sequence_parallel_value, @@ -625,7 +635,7 @@ def forward_loss( denoised_timestep_from=denoised_timestep_from, denoised_timestep_to=denoised_timestep_to, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) noise = self.student.random_noise_like( generated, @@ -753,7 +763,7 @@ def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): sigma = self.scheduler.sigma_at( idx, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) context = torch.enable_grad if (grad_enabled and idx == end_step_idx) else torch.no_grad with context(): @@ -782,23 +792,19 @@ def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): return self.student.to_dtype( x0, - self.running_dtype, + self.latent_dtype, ), *self._denoised_timestep_window(end_step_idx) def _sample_score_sigma(self, denoised_timestep_from, denoised_timestep_to, device, dtype): - min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep - max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep - min_timestep = max(0, int(min_timestep)) - max_timestep = min(self.num_train_timestep, int(max_timestep)) - if max_timestep <= min_timestep: - max_timestep = min(self.num_train_timestep, min_timestep + 1) - - timestep = torch.randint(min_timestep, max_timestep, (1,), device=device, dtype=torch.long).float() - if self.score_timestep_shift > 1: - t = timestep / self.num_train_timestep - timestep = self.score_timestep_shift * t / (1 + (self.score_timestep_shift - 1) * t) * self.num_train_timestep - timestep = timestep.clamp(self.min_step, self.max_step) - return broadcast_sequence_parallel_value((timestep / self.num_train_timestep).to(dtype=dtype)) + sigma = self.score_sigma_sampler.sample( + ScoreSigmaContext( + denoised_timestep_from=denoised_timestep_from, + denoised_timestep_to=denoised_timestep_to, + num_train_timesteps=self.num_train_timestep, + device=device, + ) + ) + return broadcast_sequence_parallel_value(sigma).to(dtype=dtype) def _denoised_timestep_window(self, exit_idx): exit_idx = int(exit_idx) diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py index 7f0ee24d2..246a72871 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/rollout.py @@ -58,7 +58,7 @@ def _run_region_rollout( sigma = self.scheduler.sigma_at( step_index, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) keep_gradient = grad_enabled and step_index == gradient_step_index context = torch.enable_grad if keep_gradient else torch.no_grad @@ -85,7 +85,7 @@ def _run_region_rollout( ) if region == "high": sigma_s = self._phase_sigma( - dtype=self.running_dtype, + dtype=self.latent_dtype, ) anchor = self._euler_step( xt, @@ -98,8 +98,8 @@ def _run_region_rollout( anchor = student_x0 raw_gradient_timestep = self._raw_timestep_from_warped_step(self.denoising_steps[gradient_step_index]) return ( - anchor.to(dtype=self.running_dtype), - student_x0.to(dtype=self.running_dtype), + anchor.to(dtype=self.latent_dtype), + student_x0.to(dtype=self.latent_dtype), sigma_s, raw_gradient_timestep, ) @@ -127,7 +127,7 @@ def _run_full_fake_rollout( sigma = self.scheduler.sigma_at( step_index, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) velocity = self._predict_velocity( active_model, @@ -145,8 +145,8 @@ def _run_full_fake_rollout( if x_bound is None: raise RuntimeError("Full phased rollout did not reach the phase boundary.") return ( - x_bound.to(dtype=self.running_dtype), - xt.detach().to(dtype=self.running_dtype), + x_bound.to(dtype=self.latent_dtype), + xt.detach().to(dtype=self.latent_dtype), ) def _teacher_velocity( @@ -200,7 +200,7 @@ def _dmd_loss_for_region( raw_min, raw_max, anchor.device, - self.running_dtype, + self.latent_dtype, ) noise = broadcast_sequence_parallel_value(torch.randn_like(anchor, dtype=torch.float32)) fake_model = self._fake_model_for_dmd( @@ -269,7 +269,7 @@ def _fake_loss_for_score_range( raw_min, raw_max, anchor.device, - self.running_dtype, + self.latent_dtype, ) noise = broadcast_sequence_parallel_value(torch.randn_like(anchor, dtype=torch.float32)) with torch.no_grad(): @@ -302,7 +302,7 @@ def _fake_score_specifications(self, region, x_bound, x0): if region == "high": anchor = x_bound sigma_s = self._phase_sigma( - dtype=self.running_dtype, + dtype=self.latent_dtype, ) fake_model = self.fake raw_min = self.match_timestep + self.score_timestep_margin @@ -312,7 +312,7 @@ def _fake_score_specifications(self, region, x_bound, x0): sigma_s = torch.zeros( 1, device=x0.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) fake_model = self.fake_2 raw_min = 1 @@ -336,7 +336,7 @@ def _fake_score_specifications(self, region, x_bound, x0): torch.zeros( 1, device=x0.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ), (self.match_timestep + self.score_timestep_margin), self.num_train_timestep, @@ -347,7 +347,7 @@ def _fake_score_specifications(self, region, x_bound, x0): def _extract_real_latents(self, sample): return self.student.extract_real_latents( sample, - self.running_dtype, + self.latent_dtype, broadcast_sequence_parallel_value, ) @@ -396,7 +396,7 @@ def _diversity_step_context(self, initial_noise, conditions): negative_condition=conditions[1], latent_hw=initial_noise.shape[-2:], device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, predict_teacher_velocity=partial( self._teacher_velocity, self.teacher, @@ -431,7 +431,7 @@ def _real_data_fake_context(self, sample, conditions, region): condition=conditions[0], negative_condition=conditions[1], device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, extract_real_latents=self._extract_real_latents, sample_synced_int=self._sample_synced_int, broadcast_noise=broadcast_sequence_parallel_value, diff --git a/lightx2v_train/lightx2v_train/trainers/sgmd.py b/lightx2v_train/lightx2v_train/trainers/sgmd.py index 4a9ff8ba6..c1ee2b1a7 100644 --- a/lightx2v_train/lightx2v_train/trainers/sgmd.py +++ b/lightx2v_train/lightx2v_train/trainers/sgmd.py @@ -102,7 +102,7 @@ def _generator_forward( denoised_timestep_from=denoised_timestep_from, denoised_timestep_to=denoised_timestep_to, device=self.student.device, - dtype=self.running_dtype, + dtype=self.latent_dtype, ) noise = broadcast_sequence_parallel_value( torch.randn( diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index d29883fa9..7f6526388 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -96,7 +96,6 @@ def merge(self, other_register): "dmd": "lightx2v_train.trainers.dmd.trainer", "dopsd": "lightx2v_train.trainers.dopsd", "flow_matching": "lightx2v_train.trainers.flow_matching", - "minimax_h3_t2av_dmd": "lightx2v_train.trainers.dmd.minimax_h3_trainer", "phased_dmd": "lightx2v_train.trainers.phased_dmd.trainer", "sgmd": "lightx2v_train.trainers.sgmd", "teacher_forcing": "lightx2v_train.trainers.teacher_forcing", diff --git a/lightx2v_train/tests/test_dmd_score_sampling.py b/lightx2v_train/tests/test_dmd_score_sampling.py new file mode 100644 index 000000000..02760786a --- /dev/null +++ b/lightx2v_train/tests/test_dmd_score_sampling.py @@ -0,0 +1,126 @@ +import pytest +import torch +from lightx2v_train.trainers.dmd.config import DmdConfig +from lightx2v_train.trainers.dmd.score_sampling import ( + ContinuousUniformScoreSigmaSampler, + DiscreteTimestepScoreSigmaSampler, + ScoreSigmaContext, + build_score_sigma_sampler, +) + + +def _context(**overrides): + values = { + "denoised_timestep_from": None, + "denoised_timestep_to": None, + "num_train_timesteps": 1000, + "device": torch.device("cpu"), + } + values.update(overrides) + return ScoreSigmaContext(**values) + + +def test_discrete_score_sampler_preserves_rollout_bounds_shift_and_clamp(monkeypatch): + sampled_bounds = [] + + def fake_randint(low, high, size, **kwargs): + sampled_bounds.append((low, high, size)) + return torch.tensor([400], device=kwargs["device"], dtype=kwargs["dtype"]) + + monkeypatch.setattr(torch, "randint", fake_randint) + sampler = DiscreteTimestepScoreSigmaSampler( + sample_min_timestep=50, + clamp_min_timestep=20, + clamp_max_timestep=980, + timestep_shift=2.0, + use_rollout_min=True, + use_rollout_max=True, + ) + + sigma = sampler.sample( + _context( + denoised_timestep_from=900, + denoised_timestep_to=200, + ) + ) + + assert sampled_bounds == [(200, 900, (1,))] + assert sigma.item() == pytest.approx(2.0 * 0.4 / (1.0 + 0.4)) + + +def test_discrete_score_sampler_clamps_after_sampling(): + sampler = DiscreteTimestepScoreSigmaSampler( + sample_min_timestep=0, + clamp_min_timestep=999, + clamp_max_timestep=1000, + ) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + sigma = sampler.sample(_context()) + + assert sigma.item() == pytest.approx(0.999) + + +def test_continuous_score_sampler_returns_unshifted_fp32_sigma(): + sampler = ContinuousUniformScoreSigmaSampler(min_sigma=0.2, max_sigma=0.3) + + sigma = sampler.sample(_context()) + + assert sigma.shape == (1,) + assert sigma.dtype == torch.float32 + assert 0.2 <= sigma.item() < 0.3 + + +def test_score_sampler_builder_defaults_to_legacy_discrete_policy(): + sampler = build_score_sigma_sampler( + None, + sample_min_timestep=10, + clamp_min_timestep=20, + clamp_max_timestep=980, + timestep_shift=5.0, + use_rollout_min=True, + use_rollout_max=False, + ) + + assert sampler == DiscreteTimestepScoreSigmaSampler( + sample_min_timestep=10, + clamp_min_timestep=20, + clamp_max_timestep=980, + timestep_shift=5.0, + use_rollout_min=True, + use_rollout_max=False, + ) + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ("continuous_uniform", "must be a mapping"), + ({"type": "unknown"}, "Unsupported"), + ({"type": "continuous_uniform", "min_sigma": 0.8, "max_sigma": 0.2}, "bounds"), + ], +) +def test_score_sampler_builder_rejects_invalid_config(config, message): + with pytest.raises(ValueError, match=message): + build_score_sigma_sampler( + config, + sample_min_timestep=0, + clamp_min_timestep=20, + clamp_max_timestep=980, + timestep_shift=5.0, + use_rollout_min=False, + use_rollout_max=False, + ) + + +def test_dmd_config_parses_an_optional_latent_dtype(): + config = { + "training": { + "student": {"train_type": "full"}, + "fake": {"train_type": "full"}, + "teacher": {}, + "dmd": {"latent_dtype": "fp32"}, + } + } + + assert DmdConfig.from_mapping(config).latent_dtype == torch.float32 From e48e6aabbcbaca846611cca412da2fe59dca3dd9 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 18 Aug 2026 11:20:29 +0000 Subject: [PATCH 13/34] update --- .../train/dmd/minimax_h3_t2av_dmd_lora.yaml | 2 +- .../model_capabilities/__init__.py | 16 +++++----- .../model_capabilities/training.py | 14 ++++----- .../lightx2v_train/model_zoo/base.py | 2 +- .../model_zoo/capability_adapters/common.py | 24 +++++++-------- .../{consistency.py => consistency_model.py} | 4 +-- .../flux2/capability_adapters/__init__.py | 4 +-- ... => flux2_consistency_model_capability.py} | 8 ++--- .../model_zoo/flux2/flux2_dev.py | 14 ++++----- .../model_zoo/flux2/flux2_klein.py | 16 +++++----- .../capability_adapters/__init__.py | 4 +-- ...cat_image_consistency_model_capability.py} | 8 ++--- .../model_zoo/longcat_image/longcat_image.py | 14 ++++----- .../ltx/capability_adapters/__init__.py | 4 +-- ...ssive_distribution_matching_capability.py} | 28 ++++++++--------- ...py => ltx_consistency_model_capability.py} | 6 ++-- ...> ltx_distribution_matching_capability.py} | 4 +-- .../lightx2v_train/model_zoo/ltx/ltx_t2av.py | 26 ++++++++-------- .../capability_adapters/__init__.py | 8 ++--- ...inimax_h3_consistency_model_capability.py} | 6 ++-- ...ax_h3_distribution_matching_capability.py} | 18 +++++------ .../model_zoo/minimax_h3/minimax_h3_t2av.py | 18 ++++++----- .../capability_adapters/__init__.py | 4 +-- ...wen_image_consistency_model_capability.py} | 6 ++-- .../model_zoo/qwen_image/qwen_image.py | 14 ++++----- .../model_zoo/qwen_image/qwen_image_edit.py | 8 ++--- .../wan/capability_adapters/__init__.py | 6 ++-- ...ssive_distribution_matching_capability.py} | 8 ++--- ...py => wan_consistency_model_capability.py} | 10 +++---- ...> wan_distribution_matching_capability.py} | 6 ++-- .../model_zoo/wan/lingbot_video.py | 16 +++++----- .../lightx2v_train/model_zoo/wan/wan_t2v.py | 28 ++++++++--------- .../trainers/consistency/base.py | 2 +- .../trainers/consistency/trainer.py | 30 +++++++++---------- .../trainers/dmd/autoregressive_dmd.py | 6 ++-- .../lightx2v_train/trainers/dmd/runtime.py | 14 ++++----- .../lightx2v_train/trainers/dmd/trainer.py | 4 +-- .../trainers/phased_dmd/roles.py | 12 ++++---- 38 files changed, 212 insertions(+), 210 deletions(-) rename lightx2v_train/lightx2v_train/model_zoo/capability_adapters/{consistency.py => consistency_model.py} (98%) rename lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/{flux2_consistency_capability.py => flux2_consistency_model_capability.py} (68%) rename lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/{longcat_image_consistency_capability.py => longcat_image_consistency_model_capability.py} (66%) rename lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/{ltx_autoregressive_distillation_capability.py => ltx_autoregressive_distribution_matching_capability.py} (93%) rename lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/{ltx_consistency_capability.py => ltx_consistency_model_capability.py} (97%) rename lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/{ltx_distillation_capability.py => ltx_distribution_matching_capability.py} (99%) rename lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/{minimax_h3_consistency_capability.py => minimax_h3_consistency_model_capability.py} (98%) rename lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/{minimax_h3_distillation_capability.py => minimax_h3_distribution_matching_capability.py} (96%) rename lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/{qwen_image_consistency_capability.py => qwen_image_consistency_model_capability.py} (78%) rename lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/{wan_autoregressive_distillation_capability.py => wan_autoregressive_distribution_matching_capability.py} (98%) rename lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/{wan_consistency_capability.py => wan_consistency_model_capability.py} (73%) rename lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/{wan_distillation_capability.py => wan_distribution_matching_capability.py} (95%) 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 a5bfdc26e..2811f21c2 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 @@ -15,7 +15,7 @@ model: # Full SDPA is impractical at H3 sequence lengths; this matches run_h3.py. attention_backend: _flash_3_hub capabilities: - distillation: + distribution_matching: # The H3 denoiser applies separate shifts to one shared base sigma. video_flow_shift: 6.0 audio_flow_shift: 3.0 diff --git a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py index 5122678ae..ff3da19d8 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -7,12 +7,12 @@ ) from .training import ( AdapterBankCapability, - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, AutoregressiveRolloutContext, CheckpointCapability, - ConsistencyCapability, - DistillationCapability, - DistillationProfile, + ConsistencyModelCapability, + DistributionMatchingCapability, + DistributionMatchingProfile, DopsdCapability, DopsdPreparedBatch, DopsdPreparedTeacherBatch, @@ -28,16 +28,16 @@ __all__ = [ "AdapterBankCapability", - "AutoregressiveDistillationCapability", + "AutoregressiveDistributionMatchingCapability", "AutoregressiveRolloutContext", "BoundCapability", "CapabilityNotSupportedError", "CapabilityProvider", "CapabilityRegistry", "CheckpointCapability", - "ConsistencyCapability", - "DistillationCapability", - "DistillationProfile", + "ConsistencyModelCapability", + "DistributionMatchingCapability", + "DistributionMatchingProfile", "DopsdCapability", "DopsdPreparedBatch", "DopsdPreparedTeacherBatch", diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 395c82b9d..849ee6686 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -54,7 +54,7 @@ def compute_loss( """Compute a differentiable teacher-forcing loss.""" -class ConsistencyCapability(ModelCapability): +class ConsistencyModelCapability(ModelCapability): @abstractmethod def configure(self, features: Collection[str]) -> None: """Install model structure required by the selected objective.""" @@ -97,8 +97,8 @@ def denoiser(self): @dataclass(frozen=True) -class DistillationProfile: - """Feature and precision constraints published by a distillation adapter.""" +class DistributionMatchingProfile: + """Feature and precision constraints published by a distribution-matching adapter.""" supports_guidance: bool = True supports_cdm: bool = True @@ -109,12 +109,12 @@ class DistillationProfile: default_latent_dtype: torch.dtype | None = None -class DistillationCapability(ModelCapability): +class DistributionMatchingCapability(ModelCapability): """Model operations required by distribution-matching algorithms.""" @property - def profile(self) -> DistillationProfile: - return DistillationProfile() + def profile(self) -> DistributionMatchingProfile: + return DistributionMatchingProfile() @property @abstractmethod @@ -237,7 +237,7 @@ class AutoregressiveRolloutContext: grad_enabled: bool -class AutoregressiveDistillationCapability(ModelCapability): +class AutoregressiveDistributionMatchingCapability(ModelCapability): @abstractmethod def rollout( self, diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 96e5d4245..730627d4f 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -133,7 +133,7 @@ def denoiser_prediction_type(self): LightX2V's current diffusion backbones are trained as rectified-flow velocity predictors. Keeping this declaration on the model avoids - baking that assumption into distillation objectives and leaves room + baking that assumption into training objectives and leaves room for models that predict x0, noise, or another parameterization. """ return "velocity" diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py index 54b91536a..90f93c800 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -10,8 +10,8 @@ from lightx2v_train.model_capabilities import ( BoundCapability, CheckpointCapability, - ConsistencyCapability, - DistillationCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, FlowMatchingSFTCapability, LossResult, ParallelCapability, @@ -113,20 +113,20 @@ def state_module(self): class CommonCheckpointCapability(BoundCapability, CheckpointCapability): - _CONSISTENCY_AUXILIARY_WEIGHTS_NAME = "consistency_auxiliary.safetensors" + _CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME = "consistency_auxiliary.safetensors" - def _consistency_auxiliary_parameter_names(self) -> tuple[str, ...]: + def _consistency_model_auxiliary_parameter_names(self) -> tuple[str, ...]: capabilities = self.model.ensure_capabilities() - if not capabilities.supports(ConsistencyCapability): + if not capabilities.supports(ConsistencyModelCapability): return () - return capabilities.require(ConsistencyCapability).auxiliary_parameter_names() + 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_auxiliary_parameter_names(), - auxiliary_weights_name=self._CONSISTENCY_AUXILIARY_WEIGHTS_NAME, + auxiliary_parameter_names=self._consistency_model_auxiliary_parameter_names(), + auxiliary_weights_name=self._CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME, ) elif is_main_process(): torch.save( @@ -139,8 +139,8 @@ def load_weights(self, save_dir, train_type) -> None: self.model.load_lora_weights_for_resume(save_dir) self.model.load_auxiliary_weights( save_dir, - self._consistency_auxiliary_parameter_names(), - weights_name=self._CONSISTENCY_AUXILIARY_WEIGHTS_NAME, + self._consistency_model_auxiliary_parameter_names(), + weights_name=self._CONSISTENCY_MODEL_AUXILIARY_WEIGHTS_NAME, ) return path = os.path.join(save_dir, "model_state.pt") @@ -198,7 +198,7 @@ def compute_loss( return LossResult(loss=loss) -class GenericConsistencyCapability(BoundCapability, ConsistencyCapability): +class GenericConsistencyModelCapability(BoundCapability, ConsistencyModelCapability): def configure(self, features: Collection[str]) -> None: features = frozenset(features) if features: @@ -252,7 +252,7 @@ def denoiser(self): return self.model.denoiser_module() -class GenericDistillationCapability(BoundCapability, DistillationCapability): +class GenericDistributionMatchingCapability(BoundCapability, DistributionMatchingCapability): """Distribution-matching operations shared by image flow models.""" def __init__( diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py similarity index 98% rename from lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py rename to lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py index 34ec515ad..382503435 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/consistency_model.py @@ -11,7 +11,7 @@ import torch from torch import Tensor, nn -from .common import GenericConsistencyCapability +from .common import GenericConsistencyModelCapability def _resolve_attribute(root, path: str): @@ -163,7 +163,7 @@ def _sinusoidal_embedding(dimension: int, position: Tensor) -> Tensor: return torch.cat((torch.cos(phases), torch.sin(phases)), dim=1) -class TimeConditionedConsistencyCapability(GenericConsistencyCapability): +class TimeConditionedConsistencyModelCapability(GenericConsistencyModelCapability): """Generic CM adapter with optional endpoint conditioning and loss variance. The capability owns objective-specific state and modules. The model wrapper 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 index 783f85be4..e8442a789 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -1,6 +1,6 @@ """Capabilities implemented specifically by Flux2 models.""" -from .flux2_consistency_capability import Flux2ConsistencyCapability +from .flux2_consistency_model_capability import Flux2ConsistencyModelCapability from .flux2_dopsd_capability import Flux2DopsdCapability -__all__ = ["Flux2ConsistencyCapability", "Flux2DopsdCapability"] +__all__ = ["Flux2ConsistencyModelCapability", "Flux2DopsdCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py similarity index 68% rename from lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py index 3e4a42123..faf3281db 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_consistency_model_capability.py @@ -1,12 +1,12 @@ -"""Consistency capability for Flux2 backbones.""" +"""Consistency-model capability for Flux2 backbones.""" -from lightx2v_train.model_zoo.capability_adapters.consistency import ( +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( ProjectedTimeEmbeddingAdapter, - TimeConditionedConsistencyCapability, + TimeConditionedConsistencyModelCapability, ) -class Flux2ConsistencyCapability(TimeConditionedConsistencyCapability): +class Flux2ConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Bind generic consistency extensions to Flux2 timestep conditioning.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index df1c9f626..dbc11ae00 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -4,10 +4,10 @@ from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +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 GenericDistillationCapability, GenericFlowMatchingCapability -from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyCapability +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 ..base import BaseModel @@ -32,8 +32,8 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - GenericDistillationCapability( + DistributionMatchingCapability, + GenericDistributionMatchingCapability( self, latent_geometry=SpatialLatentGeometry( channels_path="transformer.config.in_channels", @@ -43,8 +43,8 @@ def register_capabilities(self): ), ) self.capabilities.register( - ConsistencyCapability, - Flux2ConsistencyCapability(self), + ConsistencyModelCapability, + Flux2ConsistencyModelCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 22ec01ac7..c49d6bff2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -6,8 +6,8 @@ from lightx2v_train.model_capabilities import ( AdapterBankCapability, - ConsistencyCapability, - DistillationCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, DopsdCapability, FlowMatchingSFTCapability, ) @@ -15,9 +15,9 @@ PeftAdapterBankCapability, SpatialLatentGeometry, ) -from lightx2v_train.model_zoo.capability_adapters.common import GenericDistillationCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability from lightx2v_train.model_zoo.flux2.capability_adapters import ( - Flux2ConsistencyCapability, + Flux2ConsistencyModelCapability, Flux2DopsdCapability, ) from lightx2v_train.utils.registry import MODEL_REGISTER @@ -44,8 +44,8 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - GenericDistillationCapability( + DistributionMatchingCapability, + GenericDistributionMatchingCapability( self, latent_geometry=SpatialLatentGeometry( channels_path="transformer.config.in_channels", @@ -54,8 +54,8 @@ def register_capabilities(self): ), ) self.capabilities.register( - ConsistencyCapability, - Flux2ConsistencyCapability(self), + ConsistencyModelCapability, + Flux2ConsistencyModelCapability(self), ) self.capabilities.register( DopsdCapability, 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 index 90927ea19..585c6d6c5 100644 --- 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 @@ -1,5 +1,5 @@ """Capabilities implemented specifically by LongCat-Image.""" -from .longcat_image_consistency_capability import LongCatImageConsistencyCapability +from .longcat_image_consistency_model_capability import LongCatImageConsistencyModelCapability -__all__ = ["LongCatImageConsistencyCapability"] +__all__ = ["LongCatImageConsistencyModelCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py similarity index 66% rename from lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py index 60549250c..2ea350ef5 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_consistency_model_capability.py @@ -1,12 +1,12 @@ -"""Consistency capability for LongCat-Image.""" +"""Consistency-model capability for LongCat-Image.""" -from lightx2v_train.model_zoo.capability_adapters.consistency import ( +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( ProjectedTimeEmbeddingAdapter, - TimeConditionedConsistencyCapability, + TimeConditionedConsistencyModelCapability, ) -class LongCatImageConsistencyCapability(TimeConditionedConsistencyCapability): +class LongCatImageConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Bind generic consistency extensions to LongCat's time embedding.""" def __init__(self, model) -> None: 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 index 617c09634..f55d48103 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -6,10 +6,10 @@ from diffusers.models.transformers import LongCatImageTransformer2DModel from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +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 GenericDistillationCapability, GenericFlowMatchingCapability -from lightx2v_train.model_zoo.longcat_image.capability_adapters import LongCatImageConsistencyCapability +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 @@ -34,8 +34,8 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - GenericDistillationCapability( + DistributionMatchingCapability, + GenericDistributionMatchingCapability( self, latent_geometry=SpatialLatentGeometry( channels_path="vae.config.latent_channels", @@ -44,8 +44,8 @@ def register_capabilities(self): ), ) self.capabilities.register( - ConsistencyCapability, - LongCatImageConsistencyCapability(self), + ConsistencyModelCapability, + LongCatImageConsistencyModelCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py index bffaba5a2..3ad2b5540 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py @@ -1,5 +1,5 @@ """Capabilities implemented specifically by LTX models.""" -from .ltx_consistency_capability import LTXConsistencyCapability +from .ltx_consistency_model_capability import LTXConsistencyModelCapability -__all__ = ["LTXConsistencyCapability"] +__all__ = ["LTXConsistencyModelCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py similarity index 93% rename from lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py index 2396fac06..ca85b2d59 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py @@ -1,4 +1,4 @@ -"""Autoregressive distillation capability for LTX models.""" +"""Autoregressive distribution-matching capability for LTX models.""" from __future__ import annotations @@ -7,25 +7,25 @@ import torch from lightx2v_train.model_capabilities import ( - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, AutoregressiveRolloutContext, BoundCapability, ) from lightx2v_train.model_zoo.native.ltx2 import Modality from .common import LTXJointLatents -from .ltx_distillation_capability import LTXDistillationCapability +from .ltx_distribution_matching_capability import LTXDistributionMatchingCapability -class LTXAutoregressiveDistillationCapability( +class LTXAutoregressiveDistributionMatchingCapability( BoundCapability, - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, ): """Cached chunk-wise rollout for LTX joint video/audio states.""" def __init__(self, model) -> None: super().__init__(model) - self.distillation = LTXDistillationCapability(model) + self.distribution_matching = LTXDistributionMatchingCapability(model) def rollout( self, @@ -56,8 +56,8 @@ def rollout( len(context.denoising_steps), latents.video.device, ) - video_positions = self.distillation._video_token_positions(latent_shape) - audio_positions = self.distillation._audio_token_positions(latent_shape) + video_positions = self.distribution_matching._video_token_positions(latent_shape) + audio_positions = self.distribution_matching._audio_token_positions(latent_shape) cache = self._new_caches( dtype=latents.video.dtype, device=latents.video.device, @@ -110,7 +110,7 @@ def rollout( video_start, audio_start, ) - x0 = self.distillation.x0_from_velocity( + x0 = self.distribution_matching.x0_from_velocity( block, velocity, sigma, @@ -122,12 +122,12 @@ def rollout( dtype=context.running_dtype, ) with torch.no_grad(): - noise = self.distillation.random_noise_like( + noise = self.distribution_matching.random_noise_like( x0, x0.video.dtype, lambda value: value, ) - block = self.distillation.add_noise( + block = self.distribution_matching.add_noise( context.trajectory_scheduler, x0, noise, @@ -136,7 +136,7 @@ def rollout( video_chunks.append(x0.video) audio_chunks.append(x0.audio) - cache_latents = self.distillation.detach(x0) + cache_latents = self.distribution_matching.detach(x0) cache_sigma = torch.full( (1,), context.context_noise, @@ -144,10 +144,10 @@ def rollout( dtype=context.running_dtype, ) if context.context_noise > 0: - cache_latents = self.distillation.add_noise( + cache_latents = self.distribution_matching.add_noise( context.trajectory_scheduler, cache_latents, - self.distillation.random_noise_like( + self.distribution_matching.random_noise_like( cache_latents, cache_latents.video.dtype, lambda value: value, diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py similarity index 97% rename from lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py index 403c37a74..caea3b7ce 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py @@ -7,8 +7,8 @@ import torch from torch import Tensor, nn -from lightx2v_train.model_zoo.capability_adapters.consistency import ( - TimeConditionedConsistencyCapability, +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( + TimeConditionedConsistencyModelCapability, TimeEmbeddingAdapter, ) from lightx2v_train.model_zoo.native.ltx2 import Modality @@ -82,7 +82,7 @@ def _add_broadcast(value: Tensor, endpoint: Tensor) -> Tensor: return value + endpoint.to(device=value.device, dtype=value.dtype) -class LTXConsistencyCapability(TimeConditionedConsistencyCapability): +class LTXConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Expose joint LTX audio/video denoising through the tensor CM boundary.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py similarity index 99% rename from lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py index 1b5810306..6c7df2a7a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py @@ -5,7 +5,7 @@ import torch from lightx2v_train.model_zoo.capability_adapters.common import ( - GenericDistillationCapability, + GenericDistributionMatchingCapability, _negative_prompt, _require_single_prompt, _require_singleton_tensor, @@ -24,7 +24,7 @@ from .common import LTXJointLatentShape, LTXJointLatents -class LTXDistillationCapability(GenericDistillationCapability): +class LTXDistributionMatchingCapability(GenericDistributionMatchingCapability): """Distribution-matching operations for joint LTX video/audio states.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py index 064499d63..f9b0a7a7a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py @@ -10,17 +10,17 @@ from safetensors.torch import load_file from lightx2v_train.model_capabilities import ( - AutoregressiveDistillationCapability, - ConsistencyCapability, - DistillationCapability, + AutoregressiveDistributionMatchingCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, FlowMatchingSFTCapability, TeacherForcingCapability, ) -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_autoregressive_distillation_capability import ( - LTXAutoregressiveDistillationCapability, +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_autoregressive_distribution_matching_capability import ( + LTXAutoregressiveDistributionMatchingCapability, ) -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_consistency_capability import LTXConsistencyCapability -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_distillation_capability import LTXDistillationCapability +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_consistency_model_capability import LTXConsistencyModelCapability +from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_distribution_matching_capability import LTXDistributionMatchingCapability from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_flow_matching_capability import LTXFlowMatchingCapability from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_teacher_forcing_capability import LTXTeacherForcingCapability from lightx2v_train.model_zoo.native.ltx2 import ( @@ -54,17 +54,17 @@ def register_capabilities(self): LTXFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - LTXDistillationCapability(self), + DistributionMatchingCapability, + LTXDistributionMatchingCapability(self), ) self.capabilities.register( - ConsistencyCapability, - LTXConsistencyCapability(self), + ConsistencyModelCapability, + LTXConsistencyModelCapability(self), ) if self.config["model"]["name"] == "ltx_t2av_ar": self.capabilities.register( - AutoregressiveDistillationCapability, - LTXAutoregressiveDistillationCapability(self), + AutoregressiveDistributionMatchingCapability, + LTXAutoregressiveDistributionMatchingCapability(self), ) self.capabilities.register( TeacherForcingCapability, 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 index 194f1b44b..dd093f5f4 100644 --- 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 @@ -1,12 +1,12 @@ """Capability adapters for MiniMax-H3.""" from .common import MiniMaxH3JointLatents, MiniMaxH3LatentShape -from .minimax_h3_consistency_capability import MiniMaxH3ConsistencyCapability -from .minimax_h3_distillation_capability import MiniMaxH3DistillationCapability +from .minimax_h3_consistency_model_capability import MiniMaxH3ConsistencyModelCapability +from .minimax_h3_distribution_matching_capability import MiniMaxH3DistributionMatchingCapability __all__ = [ - "MiniMaxH3ConsistencyCapability", - "MiniMaxH3DistillationCapability", + "MiniMaxH3ConsistencyModelCapability", + "MiniMaxH3DistributionMatchingCapability", "MiniMaxH3JointLatents", "MiniMaxH3LatentShape", ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_model_capability.py similarity index 98% rename from lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_model_capability.py index be05db679..1eb2584d4 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_consistency_model_capability.py @@ -7,8 +7,8 @@ import torch from torch import Tensor, nn -from lightx2v_train.model_zoo.capability_adapters.consistency import ( - TimeConditionedConsistencyCapability, +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( + TimeConditionedConsistencyModelCapability, TimeEmbeddingAdapter, ) from lightx2v_train.model_zoo.native.minimax_h3 import MiniMaxH3PackedSequence, build_packed_sequence @@ -38,7 +38,7 @@ def encode( return embedder(projected.to(device=parameter.device, dtype=parameter.dtype)) -class MiniMaxH3ConsistencyCapability(TimeConditionedConsistencyCapability): +class MiniMaxH3ConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Expose H3's packed joint denoiser through the consistency tensor API.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py similarity index 96% rename from lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py index d09ef3352..f8382017f 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/minimax_h3/capability_adapters/minimax_h3_distribution_matching_capability.py @@ -8,9 +8,9 @@ import torch import torch.nn.functional as F -from lightx2v_train.model_capabilities import DistillationProfile +from lightx2v_train.model_capabilities import DistributionMatchingProfile from lightx2v_train.model_zoo.capability_adapters.common import ( - GenericDistillationCapability, + GenericDistributionMatchingCapability, _require_single_prompt, _require_singleton_tensor, ) @@ -25,18 +25,18 @@ @dataclass(frozen=True) -class MiniMaxH3DistillationOptions: +class MiniMaxH3DistributionMatchingOptions: video_loss_weight: float = 1.0 audio_loss_weight: float = 1.0 video_flow_shift: float = 12.0 audio_flow_shift: float = 3.0 @classmethod - def from_mapping(cls, config: Mapping | None) -> "MiniMaxH3DistillationOptions": + def from_mapping(cls, config: Mapping | None) -> "MiniMaxH3DistributionMatchingOptions": if config is None: config = {} if not isinstance(config, Mapping): - raise ValueError("model.capabilities.distillation must be a 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)), @@ -64,7 +64,7 @@ def _expand_sigma(sigma: torch.Tensor, ndim: int) -> torch.Tensor: return sigma.reshape(sigma.shape[0], *([1] * (ndim - 1))) -class MiniMaxH3DistillationCapability(GenericDistillationCapability): +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 @@ -81,7 +81,7 @@ class MiniMaxH3DistillationCapability(GenericDistillationCapability): "ff.net.0.proj", "ff.net.2", ) - _PROFILE = DistillationProfile( + _PROFILE = DistributionMatchingProfile( supports_guidance=False, supports_cdm=False, supports_ida=False, @@ -93,7 +93,7 @@ class MiniMaxH3DistillationCapability(GenericDistillationCapability): def __init__(self, model, options: Mapping | None = None) -> None: super().__init__(model) - options = MiniMaxH3DistillationOptions.from_mapping(options) + 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 @@ -101,7 +101,7 @@ def __init__(self, model, options: Mapping | None = None) -> None: self._layout_cache = {} @property - def profile(self) -> DistillationProfile: + def profile(self) -> DistributionMatchingProfile: return self._PROFILE @property 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 index 466601230..9f79bf5d6 100644 --- 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 @@ -6,8 +6,8 @@ import torch from peft import LoraConfig, inject_adapter_in_model -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability -from lightx2v_train.model_zoo.minimax_h3.capability_adapters import MiniMaxH3ConsistencyCapability, MiniMaxH3DistillationCapability +from lightx2v_train.model_capabilities import ConsistencyModelCapability, DistributionMatchingCapability +from lightx2v_train.model_zoo.minimax_h3.capability_adapters import MiniMaxH3ConsistencyModelCapability, MiniMaxH3DistributionMatchingCapability 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 @@ -26,16 +26,18 @@ def register_capabilities(self): 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.") self.capabilities.register( - DistillationCapability, - MiniMaxH3DistillationCapability( + DistributionMatchingCapability, + MiniMaxH3DistributionMatchingCapability( self, - capability_config.get("distillation"), + capability_config.get("distribution_matching"), ), ) self.capabilities.register( - ConsistencyCapability, - MiniMaxH3ConsistencyCapability(self), + ConsistencyModelCapability, + MiniMaxH3ConsistencyModelCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): @@ -70,7 +72,7 @@ def transformer_forward_context(self): def add_lora(self, rank, alpha, target_modules): if not target_modules: - target_modules = MiniMaxH3DistillationCapability._DEFAULT_LORA_TARGETS + target_modules = MiniMaxH3DistributionMatchingCapability._DEFAULT_LORA_TARGETS lora_config = LoraConfig( r=rank, lora_alpha=alpha, diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py index a80ead5fb..35639af28 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py @@ -1,5 +1,5 @@ """Capability adapters for Qwen-Image models.""" -from .qwen_image_consistency_capability import QwenImageConsistencyCapability +from .qwen_image_consistency_model_capability import QwenImageConsistencyModelCapability -__all__ = ["QwenImageConsistencyCapability"] +__all__ = ["QwenImageConsistencyModelCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_model_capability.py similarity index 78% rename from lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_model_capability.py index e2a6ac018..ec98dc422 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_consistency_model_capability.py @@ -1,12 +1,12 @@ """Consistency-model capability for Qwen-Image backbones.""" -from lightx2v_train.model_zoo.capability_adapters.consistency import ( +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( ProjectedTimeEmbeddingAdapter, - TimeConditionedConsistencyCapability, + TimeConditionedConsistencyModelCapability, ) -class QwenImageConsistencyCapability(TimeConditionedConsistencyCapability): +class QwenImageConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Bind the generic consistency extensions to Qwen's time embedding.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 9f3f20140..6ae4740c1 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -4,10 +4,10 @@ from diffusers import AutoencoderKLQwenImage, QwenImagePipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +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 GenericDistillationCapability, GenericFlowMatchingCapability -from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyCapability +from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability +from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyModelCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -37,8 +37,8 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - GenericDistillationCapability( + DistributionMatchingCapability, + GenericDistributionMatchingCapability( self, latent_geometry=SpatialLatentGeometry( channels_path="vae.config.z_dim", @@ -47,8 +47,8 @@ def register_capabilities(self): ), ) self.capabilities.register( - ConsistencyCapability, - QwenImageConsistencyCapability(self), + ConsistencyModelCapability, + QwenImageConsistencyModelCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index 1f95206cf..e9de8e90b 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -6,9 +6,9 @@ from diffusers import AutoencoderKLQwenImage, QwenImageEditPlusPipeline, QwenImageTransformer2DModel from diffusers.image_processor import VaeImageProcessor -from lightx2v_train.model_capabilities import ConsistencyCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyModelCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability -from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyCapability +from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyModelCapability from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel @@ -52,8 +52,8 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - ConsistencyCapability, - QwenImageConsistencyCapability(self), + ConsistencyModelCapability, + QwenImageConsistencyModelCapability(self), ) def load_components(self, transformer_only=False, reference_model=None): diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py index 8b32bda2b..1555ead80 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/__init__.py @@ -1,8 +1,8 @@ """Capabilities implemented specifically by Wan-family models.""" -from .wan_consistency_capability import LingBotConsistencyCapability, WanConsistencyCapability +from .wan_consistency_model_capability import LingBotConsistencyModelCapability, WanConsistencyModelCapability __all__ = [ - "LingBotConsistencyCapability", - "WanConsistencyCapability", + "LingBotConsistencyModelCapability", + "WanConsistencyModelCapability", ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distribution_matching_capability.py similarity index 98% rename from lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distribution_matching_capability.py index 75de6798b..f9b31fb60 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_autoregressive_distribution_matching_capability.py @@ -1,4 +1,4 @@ -"""Autoregressive distillation capability for causal Wan models.""" +"""Autoregressive distribution-matching capability for causal Wan models.""" from __future__ import annotations @@ -6,7 +6,7 @@ import torch.distributed as dist from lightx2v_train.model_capabilities import ( - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, AutoregressiveRolloutContext, BoundCapability, ) @@ -18,9 +18,9 @@ from lightx2v_train.runtime.sequence_parallel import all_gather_sequence -class WanAutoregressiveDistillationCapability( +class WanAutoregressiveDistributionMatchingCapability( BoundCapability, - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, ): """Cached chunk-wise rollout for causal Wan transformers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_model_capability.py similarity index 73% rename from lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_model_capability.py index b379502f3..57129ce9a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_consistency_model_capability.py @@ -1,13 +1,13 @@ -"""Consistency capabilities for Wan-family video backbones.""" +"""Consistency-model capabilities for Wan-family video backbones.""" -from lightx2v_train.model_zoo.capability_adapters.consistency import ( +from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( ProjectedTimeEmbeddingAdapter, SinusoidalTimeEmbeddingAdapter, - TimeConditionedConsistencyCapability, + TimeConditionedConsistencyModelCapability, ) -class WanConsistencyCapability(TimeConditionedConsistencyCapability): +class WanConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Bind generic consistency extensions to Wan's sinusoidal time MLP.""" def __init__(self, model) -> None: @@ -22,7 +22,7 @@ def __init__(self, model) -> None: ) -class LingBotConsistencyCapability(TimeConditionedConsistencyCapability): +class LingBotConsistencyModelCapability(TimeConditionedConsistencyModelCapability): """Bind generic consistency extensions to LingBot's time embedding.""" def __init__(self, model) -> None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distribution_matching_capability.py similarity index 95% rename from lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py rename to lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distribution_matching_capability.py index faf3623f4..8d6f16653 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distillation_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_distribution_matching_capability.py @@ -3,7 +3,7 @@ import torch from lightx2v_train.model_zoo.capability_adapters.common import ( - GenericDistillationCapability, + GenericDistributionMatchingCapability, _negative_prompt, _require_single_prompt, _require_singleton_tensor, @@ -14,7 +14,7 @@ ) -class WanDistillationCapability(GenericDistillationCapability): +class WanDistributionMatchingCapability(GenericDistributionMatchingCapability): """Distribution-matching operations for Wan-family video models.""" @property @@ -107,7 +107,7 @@ def prepare_cached_condition(self, condition): return {"prompt_embed": prompt_embed} -class LingBotDistillationCapability(WanDistillationCapability): +class LingBotDistributionMatchingCapability(WanDistributionMatchingCapability): @property def default_negative_prompt(self): return LINGBOT_VIDEO_NEGATIVE_PROMPT diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index 0cce54550..12a07382f 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -9,11 +9,11 @@ from peft.utils import set_peft_model_state_dict from safetensors.torch import load_file -from lightx2v_train.model_capabilities import ConsistencyCapability, DistillationCapability, FlowMatchingSFTCapability +from lightx2v_train.model_capabilities import ConsistencyModelCapability, DistributionMatchingCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability -from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_capability import LingBotConsistencyCapability -from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( - LingBotDistillationCapability, +from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_model_capability import LingBotConsistencyModelCapability +from lightx2v_train.model_zoo.wan.capability_adapters.wan_distribution_matching_capability import ( + LingBotDistributionMatchingCapability, ) from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype @@ -53,12 +53,12 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - LingBotDistillationCapability(self), + DistributionMatchingCapability, + LingBotDistributionMatchingCapability(self), ) self.capabilities.register( - ConsistencyCapability, - LingBotConsistencyCapability(self), + ConsistencyModelCapability, + LingBotConsistencyModelCapability(self), ) vae_scale_factor_temporal = 4 diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py index 9d1623011..ae1171e1e 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py @@ -12,19 +12,19 @@ from safetensors.torch import load_file from lightx2v_train.model_capabilities import ( - AutoregressiveDistillationCapability, - ConsistencyCapability, - DistillationCapability, + AutoregressiveDistributionMatchingCapability, + ConsistencyModelCapability, + DistributionMatchingCapability, FlowMatchingSFTCapability, TeacherForcingCapability, ) from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability -from lightx2v_train.model_zoo.wan.capability_adapters.wan_autoregressive_distillation_capability import ( - WanAutoregressiveDistillationCapability, +from lightx2v_train.model_zoo.wan.capability_adapters.wan_autoregressive_distribution_matching_capability import ( + WanAutoregressiveDistributionMatchingCapability, ) -from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_capability import WanConsistencyCapability -from lightx2v_train.model_zoo.wan.capability_adapters.wan_distillation_capability import ( - WanDistillationCapability, +from lightx2v_train.model_zoo.wan.capability_adapters.wan_consistency_model_capability import WanConsistencyModelCapability +from lightx2v_train.model_zoo.wan.capability_adapters.wan_distribution_matching_capability import ( + WanDistributionMatchingCapability, ) from lightx2v_train.model_zoo.wan.capability_adapters.wan_teacher_forcing_capability import ( WanTeacherForcingCapability, @@ -59,17 +59,17 @@ def register_capabilities(self): GenericFlowMatchingCapability(self), ) self.capabilities.register( - DistillationCapability, - WanDistillationCapability(self), + DistributionMatchingCapability, + WanDistributionMatchingCapability(self), ) self.capabilities.register( - ConsistencyCapability, - WanConsistencyCapability(self), + ConsistencyModelCapability, + WanConsistencyModelCapability(self), ) if self.use_causal_transformer: self.capabilities.register( - AutoregressiveDistillationCapability, - WanAutoregressiveDistillationCapability(self), + AutoregressiveDistributionMatchingCapability, + WanAutoregressiveDistributionMatchingCapability(self), ) self.capabilities.register( TeacherForcingCapability, diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/base.py b/lightx2v_train/lightx2v_train/trainers/consistency/base.py index 38fa4a111..606648b8c 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/base.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/base.py @@ -93,7 +93,7 @@ class DenoiserRequest: class CapabilityDenoiser: - """Adapt a consistency capability to the objective interface.""" + """Adapt a consistency-model capability to the objective interface.""" def __init__(self, capability, path: RectifiedFlowPath): self.capability = capability diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py index 7ed329dc0..9c0365632 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py @@ -7,7 +7,7 @@ from loguru import logger from lightx2v_train.model_capabilities import ( - ConsistencyCapability, + ConsistencyModelCapability, LossResult, ParallelCapability, TrainableModelCapability, @@ -40,7 +40,7 @@ class ConsistencyTrainer(FlowMatchingTrainer): trainer_name = "consistency" required_capabilities = ( *BaseTrainer.required_capabilities, - ConsistencyCapability, + ConsistencyModelCapability, ) def __init__(self, config): @@ -55,20 +55,20 @@ def __init__(self, config): def set_model(self, model): BaseTrainer.set_model(self, model) - self.consistency = model.capabilities.require(ConsistencyCapability) + self.consistency_model = model.capabilities.require(ConsistencyModelCapability) def setup(self, resume_ckpt_path=None): self._setup_resume_checkpoint = resume_ckpt_path super().setup(resume_ckpt_path=resume_ckpt_path) self.student_denoiser = CapabilityDenoiser( - self.consistency, + self.consistency_model, self.path, ) if self.objective.requires_teacher: self.teacher_model = self._build_frozen_teacher() self.teacher_denoiser = CapabilityDenoiser( - self.teacher_consistency, + self.teacher_consistency_model, self.path, ) @@ -88,7 +88,7 @@ def setup(self, resume_ckpt_path=None): ) def _setup_trainable_model(self, model): - capability = model.ensure_capabilities().require(ConsistencyCapability) + capability = model.ensure_capabilities().require(ConsistencyModelCapability) capability.configure(self.objective.model_capabilities) super()._setup_trainable_model(model) capability.restore_trainable_auxiliary() @@ -102,7 +102,7 @@ def _setup_trainable_model(self, model): def _restore_trainable_model(self, model): super()._restore_trainable_model(model) - model.ensure_capabilities().require(ConsistencyCapability).restore_trainable_auxiliary() + model.ensure_capabilities().require(ConsistencyModelCapability).restore_trainable_auxiliary() def _load_initial_model_weights(self, model, checkpoint, *, role="student"): if not os.path.isdir(checkpoint): @@ -137,10 +137,10 @@ def _build_frozen_teacher(self): transformer_only=True, reference_model=self.model, ) - self.teacher_consistency = teacher_model.capabilities.require(ConsistencyCapability) - self.teacher_consistency.set_frozen() + self.teacher_consistency_model = teacher_model.capabilities.require(ConsistencyModelCapability) + self.teacher_consistency_model.set_frozen() teacher_model.capabilities.require(ParallelCapability).apply(self.config) - self.teacher_consistency.set_frozen() + self.teacher_consistency_model.set_frozen() logger.info( "[train] consistency teacher model={} path={}", teacher_config["model"]["name"], @@ -175,7 +175,7 @@ def _build_frozen_reference(self, spec): transformer_only=True, reference_model=self.model, ) - capability = reference_model.capabilities.require(ConsistencyCapability) + capability = reference_model.capabilities.require(ConsistencyModelCapability) if self.train_type == "lora": reference_model.capabilities.require(TrainableModelCapability).configure( "lora", @@ -202,9 +202,9 @@ def _build_frozen_reference(self, spec): def compute_loss_on_sample(self, sample): with torch.no_grad(): - clean = self.consistency.encode_latent(sample) + clean = self.consistency_model.encode_latent(sample) clean = broadcast_sequence_parallel_value(clean) - condition = self.consistency.encode_condition(sample) + condition = self.consistency_model.encode_condition(sample) condition = broadcast_sequence_parallel_value(condition) negative_condition = None @@ -215,7 +215,7 @@ def compute_loss_on_sample(self, sample): context = ConsistencyStepContext( iteration=int(getattr(self, "current_train_iteration", 0)), global_batch_size=(clean.shape[0] * get_data_parallel_world_size() * self.gradient_accumulation_iters), - latent_hw=self.consistency.sampling_latent_hw(sample, clean), + latent_hw=self.consistency_model.sampling_latent_hw(sample, clean), ) training_state = self.objective.sample_training_state( clean, @@ -263,4 +263,4 @@ def _encode_negative_condition(self, sample): else: negative_sample["conditioning"]["positive"] = cached_negative negative_sample["conditioning"]["prompt"] = encoded_prompt - return self.consistency.encode_condition(negative_sample) + return self.consistency_model.encode_condition(negative_sample) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py index df9dbda48..d5c17d0a8 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/autoregressive_dmd.py @@ -1,5 +1,5 @@ from lightx2v_train.model_capabilities import ( - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, AutoregressiveRolloutContext, ) from lightx2v_train.utils.registry import TRAINER_REGISTER @@ -12,7 +12,7 @@ class AutoregressiveDmdTrainer(DmdTrainer): trainer_name = "autoregressive_dmd" required_capabilities = ( *DmdTrainer.required_capabilities, - AutoregressiveDistillationCapability, + AutoregressiveDistributionMatchingCapability, ) supports_diversity_loss = False supports_real_data_fake = False @@ -31,7 +31,7 @@ def __init__(self, config): def set_model(self, model): super().set_model(model) - self.autoregressive = model.capabilities.require(AutoregressiveDistillationCapability) + self.autoregressive = model.capabilities.require(AutoregressiveDistributionMatchingCapability) def run_back_simulation( self, diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py index 9a56e618a..8c25a8312 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py @@ -6,7 +6,7 @@ from lightx2v_train.model_capabilities import ( CheckpointCapability, - DistillationCapability, + DistributionMatchingCapability, ParallelCapability, TrainableModelCapability, ) @@ -41,7 +41,7 @@ class _DmdRuntime(BaseTrainer): trainer_name = "dmd" required_capabilities = ( *BaseTrainer.required_capabilities, - DistillationCapability, + DistributionMatchingCapability, ) default_negative_prompt = None supports_cdm = True @@ -104,11 +104,11 @@ def __init__(self, config): def set_model(self, model): super().set_model(model) - self.student = model.capabilities.require(DistillationCapability) + self.student = model.capabilities.require(DistributionMatchingCapability) profile = self.student.profile if self._configured_latent_dtype is None and profile.default_latent_dtype is not None: self.latent_dtype = profile.default_latent_dtype - self._validate_distillation_profile(profile) + self._validate_distribution_matching_profile(profile) if self.negative_prompt is None: self.negative_prompt = self.student.default_negative_prompt for lora_config in ( @@ -120,7 +120,7 @@ def set_model(self, model): logger.info("[train] dmd latent_dtype={}", self.latent_dtype) - def _validate_distillation_profile(self, profile): + def _validate_distribution_matching_profile(self, profile): model_name = self.model_config.get("name", type(self.model).__name__) requested_features = ( (self.cdm_trick.enabled, profile.supports_cdm, "CDM"), @@ -203,7 +203,7 @@ def setup(self, resume_ckpt_path=None): transformer_only=True, reference_model=self.model, ) - self.fake = self.fake_model.capabilities.require(DistillationCapability) + self.fake = self.fake_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model(self.fake_model, role="fake") self.fake_model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: @@ -220,7 +220,7 @@ def setup(self, resume_ckpt_path=None): transformer_only=True, reference_model=self.model, ) - self.teacher = self.teacher_model.capabilities.require(DistillationCapability) + self.teacher = self.teacher_model.capabilities.require(DistributionMatchingCapability) self.teacher.denoiser().requires_grad_(False) self.teacher.set_training(False) self.teacher_model.capabilities.require(ParallelCapability).apply(self.config) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py index 326432f72..7ea364b8e 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py @@ -12,7 +12,7 @@ ) from lightx2v_train.model_capabilities import ( - DistillationCapability, + DistributionMatchingCapability, ParallelCapability, TrainableModelCapability, ) @@ -173,7 +173,7 @@ def _setup_fake_real_resources(self): transformer_only=True, reference_model=self.model, ) - self.fake_real = self.fake_real_model.capabilities.require(DistillationCapability) + self.fake_real = self.fake_real_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( self.fake_real_model, role="fake_real", diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py index d7a0eaad3..722148c6d 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py @@ -3,7 +3,7 @@ from loguru import logger from lightx2v_train.model_capabilities import ( - DistillationCapability, + DistributionMatchingCapability, ParallelCapability, TrainableModelCapability, ) @@ -248,7 +248,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): transformer_only=True, reference_model=self.model, ) - self.student_2 = self.student_2_model.capabilities.require(DistillationCapability) + self.student_2 = self.student_2_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( self.student_2_model, role="student_2", @@ -263,7 +263,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): transformer_only=True, reference_model=self.model, ) - self.fake_2 = self.fake_2_model.capabilities.require(DistillationCapability) + self.fake_2 = self.fake_2_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model(self.fake_2_model, role="fake_2") self.fake_2_model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: @@ -280,7 +280,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): reference_model=self.model, ) fake_low_high_model_path = fake_low_high_model_config["model"]["pretrained_model_name_or_path"] - self.fake_low_high = self.fake_low_high_model.capabilities.require(DistillationCapability) + self.fake_low_high = self.fake_low_high_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( self.fake_low_high_model, role="fake_low_high", @@ -308,7 +308,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): if self.gradient_checkpointing: model.capabilities.require(TrainableModelCapability).enable_gradient_checkpointing() setattr(self, f"{role}_model", model) - capability = model.capabilities.require(DistillationCapability) + capability = model.capabilities.require(DistributionMatchingCapability) setattr(self, role, capability) logger.info( "[train] phased_dmd independent {} path={} train_type={}", @@ -336,7 +336,7 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): transformer_only=True, reference_model=self.model, ) - self.teacher_2 = self.teacher_2_model.capabilities.require(DistillationCapability) + self.teacher_2 = self.teacher_2_model.capabilities.require(DistributionMatchingCapability) self.teacher_2.denoiser().requires_grad_(False) self.teacher_2.set_training(False) self.teacher_2_model.capabilities.require(ParallelCapability).apply(self.config) From b4fb43949117c66b48c1c6c6e5416a906e61ef4e Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Tue, 18 Aug 2026 11:49:23 +0000 Subject: [PATCH 14/34] update --- .../tests/test_dmd_score_sampling.py | 126 ------------ lightx2v_train/tests/test_dopsd_core.py | 194 ------------------ .../tests/test_peft_adapter_bank.py | 72 ------- 3 files changed, 392 deletions(-) delete mode 100644 lightx2v_train/tests/test_dmd_score_sampling.py delete mode 100644 lightx2v_train/tests/test_dopsd_core.py delete mode 100644 lightx2v_train/tests/test_peft_adapter_bank.py diff --git a/lightx2v_train/tests/test_dmd_score_sampling.py b/lightx2v_train/tests/test_dmd_score_sampling.py deleted file mode 100644 index 02760786a..000000000 --- a/lightx2v_train/tests/test_dmd_score_sampling.py +++ /dev/null @@ -1,126 +0,0 @@ -import pytest -import torch -from lightx2v_train.trainers.dmd.config import DmdConfig -from lightx2v_train.trainers.dmd.score_sampling import ( - ContinuousUniformScoreSigmaSampler, - DiscreteTimestepScoreSigmaSampler, - ScoreSigmaContext, - build_score_sigma_sampler, -) - - -def _context(**overrides): - values = { - "denoised_timestep_from": None, - "denoised_timestep_to": None, - "num_train_timesteps": 1000, - "device": torch.device("cpu"), - } - values.update(overrides) - return ScoreSigmaContext(**values) - - -def test_discrete_score_sampler_preserves_rollout_bounds_shift_and_clamp(monkeypatch): - sampled_bounds = [] - - def fake_randint(low, high, size, **kwargs): - sampled_bounds.append((low, high, size)) - return torch.tensor([400], device=kwargs["device"], dtype=kwargs["dtype"]) - - monkeypatch.setattr(torch, "randint", fake_randint) - sampler = DiscreteTimestepScoreSigmaSampler( - sample_min_timestep=50, - clamp_min_timestep=20, - clamp_max_timestep=980, - timestep_shift=2.0, - use_rollout_min=True, - use_rollout_max=True, - ) - - sigma = sampler.sample( - _context( - denoised_timestep_from=900, - denoised_timestep_to=200, - ) - ) - - assert sampled_bounds == [(200, 900, (1,))] - assert sigma.item() == pytest.approx(2.0 * 0.4 / (1.0 + 0.4)) - - -def test_discrete_score_sampler_clamps_after_sampling(): - sampler = DiscreteTimestepScoreSigmaSampler( - sample_min_timestep=0, - clamp_min_timestep=999, - clamp_max_timestep=1000, - ) - with torch.random.fork_rng(devices=[]): - torch.manual_seed(0) - sigma = sampler.sample(_context()) - - assert sigma.item() == pytest.approx(0.999) - - -def test_continuous_score_sampler_returns_unshifted_fp32_sigma(): - sampler = ContinuousUniformScoreSigmaSampler(min_sigma=0.2, max_sigma=0.3) - - sigma = sampler.sample(_context()) - - assert sigma.shape == (1,) - assert sigma.dtype == torch.float32 - assert 0.2 <= sigma.item() < 0.3 - - -def test_score_sampler_builder_defaults_to_legacy_discrete_policy(): - sampler = build_score_sigma_sampler( - None, - sample_min_timestep=10, - clamp_min_timestep=20, - clamp_max_timestep=980, - timestep_shift=5.0, - use_rollout_min=True, - use_rollout_max=False, - ) - - assert sampler == DiscreteTimestepScoreSigmaSampler( - sample_min_timestep=10, - clamp_min_timestep=20, - clamp_max_timestep=980, - timestep_shift=5.0, - use_rollout_min=True, - use_rollout_max=False, - ) - - -@pytest.mark.parametrize( - ("config", "message"), - [ - ("continuous_uniform", "must be a mapping"), - ({"type": "unknown"}, "Unsupported"), - ({"type": "continuous_uniform", "min_sigma": 0.8, "max_sigma": 0.2}, "bounds"), - ], -) -def test_score_sampler_builder_rejects_invalid_config(config, message): - with pytest.raises(ValueError, match=message): - build_score_sigma_sampler( - config, - sample_min_timestep=0, - clamp_min_timestep=20, - clamp_max_timestep=980, - timestep_shift=5.0, - use_rollout_min=False, - use_rollout_max=False, - ) - - -def test_dmd_config_parses_an_optional_latent_dtype(): - config = { - "training": { - "student": {"train_type": "full"}, - "fake": {"train_type": "full"}, - "teacher": {}, - "dmd": {"latent_dtype": "fp32"}, - } - } - - assert DmdConfig.from_mapping(config).latent_dtype == torch.float32 diff --git a/lightx2v_train/tests/test_dopsd_core.py b/lightx2v_train/tests/test_dopsd_core.py deleted file mode 100644 index dad31a53e..000000000 --- a/lightx2v_train/tests/test_dopsd_core.py +++ /dev/null @@ -1,194 +0,0 @@ -from contextlib import contextmanager - -import pytest -import torch - -from lightx2v_train.model_capabilities import ( - AdapterBankCapability, - DopsdCapability, - DopsdPreparedBatch, - DopsdPreparedTeacherBatch, -) -from lightx2v_train.schedulers.flow_matching import RectifiedFlowMatchingScheduler -from lightx2v_train.trainers.dopsd_core import DopsdConfig, DopsdObjective - - -class FakeAdapterBank(AdapterBankCapability): - def __init__(self): - self.active = "student" - self.training = True - self.student = torch.nn.Parameter(torch.tensor(1.0)) - self.teacher = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False) - - def configure_pair(self, *args, **kwargs): - pass - - def parameters(self, adapter_name): - return (self.student,) - - @contextmanager - def activate(self, adapter_name, training=None): - previous_adapter = self.active - previous_training = self.training - self.active = adapter_name - if training is not None: - self.training = training - try: - yield - finally: - self.active = previous_adapter - self.training = previous_training - - def set_trainable(self, adapter_name): - self.active = adapter_name - - def copy(self, source_adapter, target_adapter): - self.teacher.copy_(self.student.detach()) - - def ema_update(self, source_adapter, target_adapter, decay): - self.teacher.mul_(decay).add_(self.student.detach(), alpha=1.0 - decay) - - def load(self, path, adapter_name, weights_subdir=None): - pass - - def save(self, path, adapter_name, weights_subdir=None): - pass - - -class FakeDopsdModel(DopsdCapability): - def __init__(self, adapters): - self.adapters = adapters - - @property - def device(self): - return torch.device("cpu") - - @staticmethod - def _initial_state(): - return torch.ones(1, 2, 3, 4) - - def prepare_training_batch(self, batch, teacher_prompts, running_dtype): - return DopsdPreparedBatch( - initial_state=self._initial_state(), - state_ids=torch.zeros(1, 6, 1), - student_condition="student-condition", - teacher_condition="teacher-condition", - teacher_reference="reference", - latent_hw=(3, 4), - ) - - def prepare_teacher_batch(self, reference_image, teacher_prompts, running_dtype, generator=None): - return DopsdPreparedTeacherBatch( - initial_state=self._initial_state(), - state_ids=torch.zeros(1, 6, 1), - condition="teacher-condition", - reference="reference", - latent_hw=(3, 4), - height=24, - width=32, - ) - - def predict_velocity(self, state, time, condition, state_ids, reference=None): - weight = self.adapters.student if self.adapters.active == "student" else self.adapters.teacher - return torch.ones_like(state) * weight - - def decode_state(self, state, state_ids): - return state - - -class FakeScheduler: - def build_inference_sigmas(self, num_steps, latent_hw=None): - assert num_steps == 2 - assert latent_hw == (3, 4) - return torch.tensor([1.0, 0.5, 0.0]) - - -def _config(**overrides): - values = { - "num_training_steps": 2, - "step_loss_weights": (2.0, 1.0), - } - values.update(overrides) - config = DopsdConfig(**values) - config.validate() - return config - - -def test_dopsd_objective_matches_weighted_x0_regression_and_backpropagates(): - adapters = FakeAdapterBank() - objective = DopsdObjective( - FakeDopsdModel(adapters), - adapters, - FakeScheduler(), - _config(), - torch.float32, - ) - result = objective.compute_loss( - {"conditioning": {"prompt": "a prompt"}}, - collect_trajectory=True, - ) - - # teacher velocity=0.25 and student velocity=1.0, so the x0 delta is - # sigma * 0.75 at sigmas 1.0 and 0.5. - assert result.loss.item() == pytest.approx((2.0 * 0.75**2 + 0.375**2) / 3.0) - assert len(result.student_trajectory) == 2 - assert result.state_ids is not None - result.loss.backward() - assert adapters.student.grad.item() == pytest.approx(1.125) - assert adapters.teacher.grad is None - assert adapters.active == "student" - assert adapters.training is True - - -def test_teacher_rollout_uses_the_same_schedule_and_restores_adapter_state(): - adapters = FakeAdapterBank() - objective = DopsdObjective( - FakeDopsdModel(adapters), - adapters, - FakeScheduler(), - _config(), - torch.float32, - ) - rollout = objective.rollout_teacher( - torch.zeros(1, 3, 24, 32), - ["teacher prompt"], - num_steps=2, - ) - - assert torch.allclose(rollout.state, torch.full_like(rollout.state, 0.75)) - assert adapters.active == "student" - assert adapters.training is True - - -@pytest.mark.parametrize( - ("dopsd", "message"), - [ - ({"num_training_steps": 0}, "num_training_steps"), - ({"ema_decay": 1.0}, "ema_decay"), - ({"student_adapter": "same", "teacher_adapter": "same"}, "must be different"), - ({"step_loss_weights": [0.0, 0.0]}, "at least one positive"), - ({"trajectory_every_iters": 0}, "trajectory_every_iters"), - ], -) -def test_dopsd_config_rejects_invalid_values(dopsd, message): - with pytest.raises(ValueError, match=message): - DopsdConfig.from_training_config({"train_type": "lora", "dopsd": dopsd}) - - -def test_dopsd_config_rejects_full_training(): - with pytest.raises(ValueError, match="must be 'lora'"): - DopsdConfig.from_training_config({"train_type": "full"}) - - -def test_scheduler_can_build_a_schedule_without_mutating_inference_state(): - scheduler = RectifiedFlowMatchingScheduler( - { - "model": {"running_dtype": "fp32"}, - "scheduler": {"num_train_timesteps": 1000}, - } - ) - sigmas = scheduler.build_inference_sigmas(2, latent_hw=(3, 4)) - - assert torch.allclose(sigmas.cpu(), torch.tensor([1.0, 0.5, 0.0])) - assert scheduler.infer_sigmas is None - assert scheduler.infer_timesteps is None diff --git a/lightx2v_train/tests/test_peft_adapter_bank.py b/lightx2v_train/tests/test_peft_adapter_bank.py deleted file mode 100644 index 1064f94cf..000000000 --- a/lightx2v_train/tests/test_peft_adapter_bank.py +++ /dev/null @@ -1,72 +0,0 @@ -import torch - -from lightx2v_train.model_zoo.capability_adapters import PeftAdapterBankCapability - - -class FakeDenoiser(torch.nn.Module): - def __init__(self): - super().__init__() - self.lora_A = torch.nn.ModuleDict( - { - "student": torch.nn.Linear(2, 1, bias=False), - "teacher": torch.nn.Linear(2, 1, bias=False), - } - ) - self.active_adapter = None - - def set_adapter(self, adapter_name): - self.active_adapter = adapter_name - - -class FakeModel: - def __init__(self): - self.denoiser = FakeDenoiser() - - def denoiser_module(self): - return self.denoiser - - def load_lora_weights_for_resume(self, *args, **kwargs): - pass - - def save_lora_weights(self, *args, **kwargs): - pass - - -def test_adapter_copy_ema_trainability_and_scoped_activation_are_strict(): - model = FakeModel() - adapters = PeftAdapterBankCapability(model) - student = model.denoiser.lora_A["student"].weight - teacher = model.denoiser.lora_A["teacher"].weight - student.data.fill_(2.0) - teacher.data.zero_() - - adapters.copy("student", "teacher") - assert torch.allclose(teacher, torch.full_like(teacher, 2.0)) - - student.data.fill_(4.0) - adapters.ema_update("student", "teacher", decay=0.5) - assert torch.allclose(teacher, torch.full_like(teacher, 3.0)) - - adapters.set_trainable("student") - assert student.requires_grad - assert not teacher.requires_grad - assert model.denoiser.active_adapter == "student" - - with adapters.activate("teacher", training=False): - assert model.denoiser.active_adapter == "teacher" - assert not model.denoiser.training - assert model.denoiser.active_adapter == "student" - assert model.denoiser.training - - -def test_adapter_pairing_fails_instead_of_silently_skipping_missing_weights(): - model = FakeModel() - del model.denoiser.lora_A["teacher"] - adapters = PeftAdapterBankCapability(model) - - try: - adapters.copy("student", "teacher") - except RuntimeError as error: - assert "No parameters found" in str(error) - else: - raise AssertionError("Expected strict adapter pairing to reject a missing teacher.") From a6990dae74420c843223b50bedffbe48a7fe1e41 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Fri, 21 Aug 2026 13:54:30 +0000 Subject: [PATCH 15/34] update qwen image edit --- .../capability_adapters/__init__.py | 3 +- ...e_edit_distribution_matching_capability.py | 27 ++ .../model_zoo/qwen_image/qwen_image.py | 77 ++++-- .../model_zoo/qwen_image/qwen_image_edit.py | 241 +++--------------- 4 files changed, 121 insertions(+), 227 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py index 35639af28..a2857f10a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/__init__.py @@ -1,5 +1,6 @@ """Capability adapters for Qwen-Image models.""" from .qwen_image_consistency_model_capability import QwenImageConsistencyModelCapability +from .qwen_image_edit_distribution_matching_capability import QwenImageEditDistributionMatchingCapability -__all__ = ["QwenImageConsistencyModelCapability"] +__all__ = ["QwenImageConsistencyModelCapability", "QwenImageEditDistributionMatchingCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py new file mode 100644 index 000000000..a6106bed7 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py @@ -0,0 +1,27 @@ +"""Distribution-matching capability for Qwen-Image-Edit.""" + +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _negative_prompt, + _require_single_prompt, +) + + +class QwenImageEditDistributionMatchingCapability(GenericDistributionMatchingCapability): + """Keep source-image conditioning in Qwen-Image-Edit DMD.""" + + def encode_conditions(self, batch, negative_prompt, guidance_scale, broadcast): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + 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/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 6ae4740c1..340b9ef0a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -16,6 +16,7 @@ @dataclass class QwenImageDenoiserInput: hidden_states: torch.Tensor + target_token_length: int img_shapes: list height: int width: int @@ -29,6 +30,7 @@ class QwenImageModel(BaseModel): """ pipeline_cls = QwenImagePipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability def register_capabilities(self): super().register_capabilities() @@ -38,7 +40,7 @@ def register_capabilities(self): ) self.capabilities.register( DistributionMatchingCapability, - GenericDistributionMatchingCapability( + self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( channels_path="vae.config.z_dim", @@ -61,7 +63,7 @@ def load_components(self, transformer_only=False, reference_model=None): self.transformer = self.load_transformer() return model_path = self.config["model"]["pretrained_model_name_or_path"] - self.text_pipeline = QwenImagePipeline.from_pretrained( + self.text_pipeline = self.pipeline_cls.from_pretrained( model_path, transformer=None, vae=None, @@ -79,6 +81,9 @@ def load_transformer(self): model_path = self.config["model"]["pretrained_model_name_or_path"] return QwenImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) + def load_full_weights_for_resume(self, resume_ckpt_path): + self.transformer = QwenImageTransformer2DModel.from_pretrained(resume_ckpt_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) + def denoiser_module(self): return self.transformer @@ -100,56 +105,76 @@ def fsdp2_shard_plan(self, fsdp_config): }, ] + def _latent_statistics(self, latents): + shape = (1, self.vae.config.z_dim, 1, 1, 1) + mean = latents.new_tensor(self.vae.config.latents_mean).view(shape) + std = latents.new_tensor(self.vae.config.latents_std).view(shape) + return mean, std + + def _normalize_latents(self, latents): + mean, std = self._latent_statistics(latents) + return (latents - mean) / std + + def _denormalize_latents(self, latents): + mean, std = self._latent_statistics(latents) + return latents * std + mean + def encode_to_latent(self, sample): image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) pixel_values = image.unsqueeze(2) latent = self.vae.encode(pixel_values).latent_dist.sample() # (B, C, T, H, W) - - latent_mean = torch.tensor(self.vae.config.latents_mean, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - latent_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - return (latent - latent_mean) * latent_std + return self._normalize_latents(latent) def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) - def encode_prompt_condition(self, prompt): + def encode_prompt_condition(self, prompt, **kwargs): prompt_embed, prompt_embed_mask = self.text_pipeline.encode_prompt( prompt=prompt, device=self.device, num_images_per_prompt=1, max_sequence_length=self.config["model"].get("max_sequence_length", 1024), + **kwargs, ) return { "prompt_embed": prompt_embed, "prompt_embed_mask": prompt_embed_mask, } + def _get_additional_image_tokens(self, condition): + return None, [] + def prepare_denoiser_input(self, noisy_latent, condition=None): # noisy_latent: (B, C, T, H, W) n = noisy_latent.shape[0] h, w = noisy_latent.shape[3], noisy_latent.shape[4] - packed = QwenImagePipeline._pack_latents(noisy_latent, n, noisy_latent.shape[1], h, w) + packed = self.pipeline_cls._pack_latents(noisy_latent, n, noisy_latent.shape[1], h, w) + additional_tokens, additional_shapes = self._get_additional_image_tokens(condition) + hidden_states = packed if additional_tokens is None else torch.cat([packed, additional_tokens], dim=1) return QwenImageDenoiserInput( - hidden_states=packed, - img_shapes=[(1, h // 2, w // 2)] * n, + hidden_states=hidden_states, + target_token_length=packed.shape[1], + img_shapes=[[(1, h // 2, w // 2), *additional_shapes] for _ in range(n)], height=h, width=w, ) def denoise(self, denoiser_input, timestep_or_sigma, condition): - return self.transformer( + prediction = self.transformer( hidden_states=denoiser_input.hidden_states, timestep=timestep_or_sigma, # timestep_or_sigma is in [0, 1] not [0, 1000] guidance=None, encoder_hidden_states_mask=condition["prompt_embed_mask"], encoder_hidden_states=condition["prompt_embed"], img_shapes=denoiser_input.img_shapes, + attention_kwargs={}, return_dict=False, )[0] + return prediction[:, : denoiser_input.target_token_length] def postprocess_denoiser_output(self, prediction, denoiser_input): - return QwenImagePipeline._unpack_latents( + return self.pipeline_cls._unpack_latents( prediction, height=denoiser_input.height * self.vae_scale_factor, width=denoiser_input.width * self.vae_scale_factor, @@ -163,26 +188,24 @@ def prepare_infer_latents(self, height, width, generator=None): return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) def decode_latent(self, latent): - # Reverse the normalization from encode_to_latent: - # encode: normalized = (raw - mean) / latents_std - # decode: raw = normalized * latents_std + mean - latent_mean = torch.tensor(self.vae.config.latents_mean, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - latent_std = torch.tensor(self.vae.config.latents_std, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - latent = latent * latent_std + latent_mean # (B, C, T, H, W), C == z_dim - + latent = self._denormalize_latents(latent) image = self.vae.decode(latent).sample # (B, C, T, H, W) image = image[:, :, 0, :, :] # drop temporal dim -> (B, C, H, W), T == 1 - return self.image_processor.postprocess(image, output_type="pil") + def get_pipeline_extra_components(self): + return {} + def assemble_pipeline(self, scheduler=None): - return QwenImagePipeline( - 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, # use the original scheduler for bit-exact alignment with diffusers - ).to(self.device) + components = { + "tokenizer": self.text_pipeline.tokenizer, + "text_encoder": self.text_pipeline.text_encoder, + "vae": self.vae, + "transformer": self.transformer, + "scheduler": scheduler if scheduler is not None else self.text_pipeline.scheduler, + **self.get_pipeline_extra_components(), + } + return self.pipeline_cls(**components).to(self.device) def get_pipeline_infer_kwargs(self, infer_config): # QwenImagePipeline uses `true_cfg_scale` instead of the standard `guidance_scale` diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index e9de8e90b..1e3a575e4 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -1,28 +1,16 @@ import math -from dataclasses import dataclass import torch from PIL import Image -from diffusers import AutoencoderKLQwenImage, QwenImageEditPlusPipeline, QwenImageTransformer2DModel -from diffusers.image_processor import VaeImageProcessor +from diffusers import QwenImageEditPlusPipeline -from lightx2v_train.model_capabilities import ConsistencyModelCapability, FlowMatchingSFTCapability -from lightx2v_train.model_zoo.capability_adapters.common import GenericFlowMatchingCapability -from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageConsistencyModelCapability +from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageEditDistributionMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER -from ..base import BaseModel +from .qwen_image import QwenImageModel CONDITION_IMAGE_AREA = 384 * 384 - - -@dataclass -class QwenImageEditDenoiserInput: - hidden_states: torch.Tensor - target_token_length: int - img_shapes: list - height: int - width: int +VAE_IMAGE_AREA = 1024 * 1024 def _calculate_dimensions(target_area, ratio): @@ -34,7 +22,7 @@ def _calculate_dimensions(target_area, ratio): @MODEL_REGISTER("qwen_image_edit") -class QwenImageEditModel(BaseModel): +class QwenImageEditModel(QwenImageModel): """Qwen-Image-Edit 2511 LoRA training. This follows the local Qwen-Image-Edit-2511 model's @@ -44,103 +32,22 @@ class QwenImageEditModel(BaseModel): """ pipeline_cls = QwenImageEditPlusPipeline - - def register_capabilities(self): - super().register_capabilities() - self.capabilities.register( - FlowMatchingSFTCapability, - GenericFlowMatchingCapability(self), - ) - self.capabilities.register( - ConsistencyModelCapability, - QwenImageConsistencyModelCapability(self), - ) - - 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.vae_scale_factor = reference_model.vae_scale_factor - self.latent_channels = reference_model.latent_channels - 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 = QwenImageEditPlusPipeline.from_pretrained( - model_path, - transformer=None, - vae=None, - torch_dtype=self.running_dtype, - ).to(self.device) - self.vae = AutoencoderKLQwenImage.from_pretrained(model_path, subfolder="vae").to(self.device, dtype=self.running_dtype) - self.transformer = QwenImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) - - self.text_pipeline.text_encoder.requires_grad_(False) - self.vae.requires_grad_(False) - self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) - self.latent_channels = self.vae.config.z_dim - self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - - def load_transformer(self): - model_path = self.config["model"]["pretrained_model_name_or_path"] - return QwenImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) - - def load_full_weights_for_resume(self, resume_ckpt_path): - self.transformer = QwenImageTransformer2DModel.from_pretrained(resume_ckpt_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"], - }, - { - "module": self.transformer, - "reshard_after_forward": reshard_config["root_reshard"], - }, - ] - - def _normalize_latents(self, latents): - latent_mean = torch.tensor(self.vae.config.latents_mean, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - latent_std = torch.tensor(self.vae.config.latents_std, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - return (latents - latent_mean) / latent_std - - def _denormalize_latents(self, latents): - latent_mean = torch.tensor(self.vae.config.latents_mean, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - latent_std = torch.tensor(self.vae.config.latents_std, device=self.device, dtype=self.running_dtype).view(1, self.vae.config.z_dim, 1, 1, 1) - return latents * latent_std + latent_mean - - def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) - pixel_values = image.unsqueeze(2) - latent = self.vae.encode(pixel_values).latent_dist.sample() - return self._normalize_latents(latent) + distribution_matching_capability_cls = QwenImageEditDistributionMatchingCapability def encode_condition(self, sample): + prompt = sample["conditioning"]["prompt"] + return self.encode_conditions_with_source(sample, [prompt])[0] + + def encode_conditions_with_source(self, sample, prompts): source_images = self._source_images_from_sample(sample) - condition_images = self._condition_images_from_source_tensors(source_images) - prompt_embed, prompt_embed_mask = self.text_pipeline.encode_prompt( - prompt=sample["conditioning"]["prompt"], - image=condition_images, - device=self.device, - num_images_per_prompt=1, - max_sequence_length=self.config["model"].get("max_sequence_length", 1024), - ) - condition = { - "prompt_embed": prompt_embed, - "prompt_embed_mask": prompt_embed_mask, - } - if source_images: - source_latents, source_img_shapes = self._encode_source_image_latents(source_images) - condition["source_latents"] = source_latents - condition["source_img_shapes"] = source_img_shapes - return condition + condition_images, vae_images = self._prepare_source_images(source_images) + conditions = [self.encode_prompt_condition(prompt, image=condition_images) for prompt in prompts] + if vae_images: + source_latents, source_img_shapes = self._encode_source_image_latents(vae_images) + for condition in conditions: + condition["source_latents"] = source_latents + condition["source_img_shapes"] = source_img_shapes + return conditions def _source_images_from_sample(self, sample): source_images = sample["inputs"].get("source_images") @@ -152,123 +59,59 @@ def _source_images_from_sample(self, sample): raise TypeError(f"source_images must contain tensors after collation, got {type(image)}") if image.ndim == 3: image = image.unsqueeze(0) + if image.ndim != 4 or image.shape[0] != 1: + raise ValueError(f"QwenImageEditPlusPipeline requires source images with batch_size=1, got {tuple(image.shape)}") tensors.append(image) return tensors - def _condition_images_from_source_tensors(self, source_images): + def _prepare_source_images(self, source_images): if not source_images: - return None - batch_size = source_images[0].shape[0] - if batch_size != 1: - raise NotImplementedError("QwenImageEditModel currently expects batch_size=1 for source image prompt encoding.") + return None, [] condition_images = [] + vae_images = [] for image in source_images: pil_image = self._tensor_to_pil(image[0]) - width, height = _calculate_dimensions(CONDITION_IMAGE_AREA, pil_image.size[0] / pil_image.size[1]) - condition_images.append(self.image_processor.resize(pil_image, height, width)) - return condition_images + ratio = pil_image.width / pil_image.height + condition_width, condition_height = _calculate_dimensions(CONDITION_IMAGE_AREA, ratio) + vae_width, vae_height = _calculate_dimensions(VAE_IMAGE_AREA, ratio) + condition_images.append(self.image_processor.resize(pil_image, condition_height, condition_width)) + vae_images.append(self.image_processor.preprocess(pil_image, vae_height, vae_width).unsqueeze(2)) + return condition_images, vae_images def _tensor_to_pil(self, image): image = ((image.detach().float().cpu().clamp(-1, 1) + 1.0) * 127.5).round().byte() array = image.permute(1, 2, 0).numpy() return Image.fromarray(array, mode="RGB") - def _encode_source_image_latents(self, source_images): + def _encode_source_image_latents(self, vae_images): packed_latents = [] img_shapes = [] - for image in source_images: + for image in vae_images: image = image.to(device=self.device, dtype=self.running_dtype) - pixel_values = image.unsqueeze(2) - latent = self.vae.encode(pixel_values).latent_dist.mode() + latent = self.vae.encode(image).latent_dist.mode() latent = self._normalize_latents(latent) n, c, _, h, w = latent.shape - packed_latents.append(QwenImageEditPlusPipeline._pack_latents(latent, n, c, h, w)) + packed_latents.append(self.pipeline_cls._pack_latents(latent, n, c, h, w)) img_shapes.append((1, h // 2, w // 2)) return torch.cat(packed_latents, dim=1), img_shapes - def prepare_denoiser_input(self, noisy_latent, condition=None): + def _get_additional_image_tokens(self, condition): if condition is None: raise ValueError("QwenImageEditModel.prepare_denoiser_input requires condition.") + return condition.get("source_latents"), condition.get("source_img_shapes", []) - n = noisy_latent.shape[0] - h, w = noisy_latent.shape[3], noisy_latent.shape[4] - packed = QwenImageEditPlusPipeline._pack_latents(noisy_latent, n, noisy_latent.shape[1], h, w) - source_latents = condition.get("source_latents") - source_img_shapes = condition.get("source_img_shapes", []) - if source_latents is None: - hidden_states = packed - else: - hidden_states = torch.cat([packed, source_latents], dim=1) - img_shapes = [[(1, h // 2, w // 2), *source_img_shapes]] * n - return QwenImageEditDenoiserInput( - hidden_states=hidden_states, - target_token_length=packed.shape[1], - img_shapes=img_shapes, - height=h, - width=w, - ) - - def denoise(self, denoiser_input, timestep_or_sigma, condition): - prompt_embed_mask = condition["prompt_embed_mask"] - if prompt_embed_mask is None: - txt_seq_lens = [condition["prompt_embed"].shape[1]] * condition["prompt_embed"].shape[0] - else: - txt_seq_lens = prompt_embed_mask.sum(dim=1).tolist() - prediction = self.transformer( - hidden_states=denoiser_input.hidden_states, - timestep=timestep_or_sigma, - guidance=None, - encoder_hidden_states_mask=prompt_embed_mask, - encoder_hidden_states=condition["prompt_embed"], - img_shapes=denoiser_input.img_shapes, - txt_seq_lens=txt_seq_lens, - attention_kwargs={}, - return_dict=False, - )[0] - return prediction[:, : denoiser_input.target_token_length] - - def postprocess_denoiser_output(self, prediction, denoiser_input): - return QwenImageEditPlusPipeline._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 - shape = (1, self.vae.config.z_dim, 1, latent_h, latent_w) - return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) - - def decode_latent(self, latent): - latent = self._denormalize_latents(latent) - image = self.vae.decode(latent).sample - image = image[:, :, 0, :, :] - return self.image_processor.postprocess(image, output_type="pil") - - def assemble_pipeline(self, scheduler=None): - return QwenImageEditPlusPipeline( - tokenizer=self.text_pipeline.tokenizer, - processor=self.text_pipeline.processor, - 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_extra_components(self): + return {"processor": self.text_pipeline.processor} def get_pipeline_infer_kwargs(self, infer_config): - kwargs = { - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "true_cfg_scale": infer_config.get("cfg_guidance_scale", 4.0), - } - if infer_config.get("height") is not None: - kwargs["height"] = infer_config["height"] - if infer_config.get("width") is not None: - kwargs["width"] = infer_config["width"] + kwargs = super().get_pipeline_infer_kwargs(infer_config) + if infer_config.get("height") is None: + kwargs.pop("height") + if infer_config.get("width") is None: + kwargs.pop("width") return kwargs def get_pipeline_sample_kwargs(self, sample): From 95f27b7657930dd54cc666d7bbb0a011e32e3896 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Fri, 21 Aug 2026 13:58:35 +0000 Subject: [PATCH 16/34] update --- .../model_zoo/qwen_image/qwen_image_edit.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index 1e3a575e4..cfea1899e 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -23,12 +23,8 @@ def _calculate_dimensions(target_area, ratio): @MODEL_REGISTER("qwen_image_edit") class QwenImageEditModel(QwenImageModel): - """Qwen-Image-Edit 2511 LoRA training. - - This follows the local Qwen-Image-Edit-2511 model's - QwenImageEditPlusPipeline conditioning path: source images are used by the - Qwen2.5-VL text encoder and are also encoded as additional VAE latent - tokens for the transformer. + """Supports weights from these Hugging Face repos: + - https://huggingface.co/Qwen/Qwen-Image-Edit-2511 """ pipeline_cls = QwenImageEditPlusPipeline From aa408b2cbbdfee2812daddc622d84e0f96bbb447 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sat, 22 Aug 2026 10:19:31 +0000 Subject: [PATCH 17/34] update image models --- .../lightx2v_train/data/image_dataset.py | 32 +--- lightx2v_train/lightx2v_train/infer/image.py | 14 +- .../model_zoo/flux2/data_process.py | 71 +++++++++ .../model_zoo/flux2/flux2_dev.py | 11 +- .../model_zoo/flux2/flux2_klein.py | 11 +- .../model_zoo/longcat_image/data_process.py | 71 +++++++++ .../model_zoo/longcat_image/longcat_image.py | 11 +- .../model_zoo/qwen_image/data_process.py | 140 ++++++++++++++++++ .../model_zoo/qwen_image/qwen_image.py | 11 +- .../model_zoo/qwen_image/qwen_image_edit.py | 56 +------ .../lightx2v_train/utils/image_ops.py | 47 ++++++ 11 files changed, 384 insertions(+), 91 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py create mode 100644 lightx2v_train/lightx2v_train/utils/image_ops.py diff --git a/lightx2v_train/lightx2v_train/data/image_dataset.py b/lightx2v_train/lightx2v_train/data/image_dataset.py index 8d362d32e..7ec4a04d9 100644 --- a/lightx2v_train/lightx2v_train/data/image_dataset.py +++ b/lightx2v_train/lightx2v_train/data/image_dataset.py @@ -1,5 +1,4 @@ import json -import math import random from pathlib import Path @@ -12,31 +11,15 @@ 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 without applying model-specific transforms.""" + def __init__( self, metadata_paths, - target_area=1024 * 1024, prompt_dropout_rate=0.0, ): - self.target_area = target_area self.prompt_dropout_rate = prompt_dropout_rate self.samples = [] for path in metadata_paths: @@ -104,10 +87,11 @@ 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: + array = np.array(image.convert("RGB"), dtype=np.uint8, copy=True) + return torch.from_numpy(array).permute(2, 0, 1) @DATA_REGISTER("image_dataset") @@ -115,14 +99,12 @@ def build_image_dataset(data_config_split, train_or_val="train"): 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, ) dp_world_size = get_data_parallel_world_size() diff --git a/lightx2v_train/lightx2v_train/infer/image.py b/lightx2v_train/lightx2v_train/infer/image.py index 284aae6a4..8d4f73e5a 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -13,18 +13,6 @@ 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 _load_infer_sample(self, index, prompt): @@ -82,7 +70,7 @@ 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 = self.model.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) 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..dea34d058 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -0,0 +1,71 @@ +import torch + +from lightx2v_train.utils.image_ops import ( + align_dimension, + calculate_area_dimensions, + image_tensor_to_pil, + resize_and_center_crop, +) + + +def _target_area_from_config(config): + 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 + + +class Flux2DataProcessor: + def __init__(self, image_processor, config): + self.image_processor = image_processor + self.target_area = _target_area_from_config(config) + + def process_target(self, sample): + image = image_tensor_to_pil(sample["inputs"]["target_image"]) + 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) + return self.image_processor.preprocess(image) + + 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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index dbc11ae00..8aac16bfe 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from functools import cached_property import torch from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel @@ -11,6 +12,7 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel +from .data_process import Flux2DataProcessor @dataclass @@ -25,6 +27,10 @@ class Flux2DevDenoiserInput: class Flux2DevModel(BaseModel): pipeline_cls = Flux2Pipeline + @cached_property + def data_processor(self): + return Flux2DataProcessor(self.image_processor, self.config) + def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -115,7 +121,7 @@ def _denormalize_patch_latents(self, latents): return Flux2Pipeline._unpatchify_latents(latents) def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) + image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) latent = self.vae.encode(image).latent_dist.sample() return self._normalize_patch_latents(latent) @@ -123,6 +129,9 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) + def infer_target_size(self, sample, default_height, default_width): + return self.data_processor.infer_target_size(sample, default_height, default_width) + def encode_prompt_condition(self, prompt): model_config = self.config["model"] prompt_embed, text_ids = self.text_pipeline.encode_prompt( diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index c49d6bff2..80a29c697 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from functools import cached_property import torch from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel @@ -23,6 +24,7 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel +from .data_process import Flux2DataProcessor @dataclass @@ -37,6 +39,10 @@ class Flux2KleinDenoiserInput: class Flux2KleinModel(BaseModel): pipeline_cls = Flux2KleinPipeline + @cached_property + def data_processor(self): + return Flux2DataProcessor(self.image_processor, self.config) + def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -129,7 +135,7 @@ def _denormalize_patch_latents(self, latents): return Flux2KleinPipeline._unpatchify_latents(latents) def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) + image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) latent = self.vae.encode(image).latent_dist.sample() return self._normalize_patch_latents(latent) @@ -137,6 +143,9 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) + def infer_target_size(self, sample, default_height, default_width): + return self.data_processor.infer_target_size(sample, default_height, default_width) + def encode_prompt_condition(self, prompt): model_config = self.config["model"] prompt_embed, text_ids = self.text_pipeline.encode_prompt( 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..49dbae0b1 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py @@ -0,0 +1,71 @@ +import torch + +from lightx2v_train.utils.image_ops import ( + align_dimension, + calculate_area_dimensions, + image_tensor_to_pil, + resize_and_center_crop, +) + + +def _target_area_from_config(config): + 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 + + +class LongCatImageDataProcessor: + def __init__(self, image_processor, config): + self.image_processor = image_processor + self.target_area = _target_area_from_config(config) + + def process_target(self, sample): + image = image_tensor_to_pil(sample["inputs"]["target_image"]) + 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) + return self.image_processor.preprocess(image) + + 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) 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 index f55d48103..1102ad306 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from functools import cached_property import torch from diffusers import AutoencoderKL, LongCatImagePipeline @@ -13,6 +14,7 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel +from .data_process import LongCatImageDataProcessor @dataclass @@ -27,6 +29,10 @@ class LongCatImageDenoiserInput: class LongCatImageModel(BaseModel): pipeline_cls = LongCatImagePipeline + @cached_property + def data_processor(self): + return LongCatImageDataProcessor(self.image_processor, self.config) + def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -106,7 +112,7 @@ 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) + image = self.data_processor.process_target(sample).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) @@ -116,6 +122,9 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) + def infer_target_size(self, sample, default_height, default_width): + return self.data_processor.infer_target_size(sample, default_height, default_width) + def encode_prompt_condition(self, prompt): rewrite_training = self.config.get( "enable_prompt_rewrite_training", diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py new file mode 100644 index 000000000..6ce883a19 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py @@ -0,0 +1,140 @@ +import torch + +from lightx2v_train.utils.image_ops import ( + align_dimension, + calculate_area_dimensions, + image_tensor_to_pil, + resize_and_center_crop, +) + +CONDITION_IMAGE_AREA = 384 * 384 +VAE_IMAGE_AREA = 1024 * 1024 + + +def _target_area_from_config(config): + 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 + + +class QwenImageDataProcessor: + def __init__(self, image_processor, config): + self.image_processor = image_processor + self.target_area = _target_area_from_config(config) + + def process_target(self, sample): + return self._process_target(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) + + def _process_target(self, sample, reference_tensor=None, area_multiple=None): + image = image_tensor_to_pil(sample["inputs"]["target_image"]) + reference = image if reference_tensor is None else image_tensor_to_pil(reference_tensor) + width, height = self._resolve_target_size( + sample, + reference, + self.target_area, + area_multiple, + ) + image = resize_and_center_crop(image, width, height) + return self.image_processor.preprocess(image) + + def _resolve_target_size(self, sample, reference_image, target_area, area_multiple=None): + height, width = _explicit_target_size(sample) + multiple = int(self.image_processor.config.vae_scale_factor) + if height is not None: + return align_dimension(width, multiple), align_dimension(height, multiple) + ratio = reference_image.width / reference_image.height + return calculate_area_dimensions(target_area, ratio, area_multiple or multiple) + + +class QwenImageEditDataProcessor(QwenImageDataProcessor): + def process_target(self, sample): + source_images = sample["inputs"].get("source_images") + reference_tensor = source_images[-1] if source_images else None + return self._process_target( + sample, + reference_tensor=reference_tensor, + area_multiple=32, + ) + + def infer_target_size(self, sample, default_height, default_width): + source_images = sample["inputs"].get("source_images") + if not source_images: + return super().infer_target_size(sample, default_height, default_width) + reference_image = image_tensor_to_pil(source_images[-1]) + width, height = self._resolve_target_size( + sample, + reference_image, + int(default_height) * int(default_width), + area_multiple=32, + ) + return height, width + + def process_source_images(self, sample): + source_images = self._source_images(sample) + if not source_images: + return None, [] + + condition_images = [] + vae_images = [] + for image in source_images: + pil_image = image_tensor_to_pil(image) + ratio = pil_image.width / pil_image.height + condition_width, condition_height = calculate_area_dimensions(CONDITION_IMAGE_AREA, ratio, multiple=32) + vae_width, vae_height = calculate_area_dimensions(VAE_IMAGE_AREA, ratio, multiple=32) + condition_images.append(self.image_processor.resize(pil_image, condition_height, condition_width)) + vae_images.append(self.image_processor.preprocess(pil_image, vae_height, vae_width).unsqueeze(2)) + return condition_images, vae_images + + @staticmethod + def _source_images(sample): + source_images = sample["inputs"].get("source_images") + if source_images is None: + return [] + tensors = [] + for image in source_images: + if not isinstance(image, torch.Tensor): + raise TypeError(f"source_images must contain tensors after collation, got {type(image)}") + if image.ndim == 3: + image = image.unsqueeze(0) + if image.ndim != 4 or image.shape[0] != 1: + raise ValueError(f"QwenImageEditPlusPipeline requires source images with batch_size=1, got {tuple(image.shape)}") + tensors.append(image) + return tensors diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 340b9ef0a..2adaf98eb 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from functools import cached_property import torch from diffusers import AutoencoderKLQwenImage, QwenImagePipeline, QwenImageTransformer2DModel @@ -11,6 +12,7 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel +from .data_process import QwenImageDataProcessor @dataclass @@ -32,6 +34,10 @@ class QwenImageModel(BaseModel): pipeline_cls = QwenImagePipeline distribution_matching_capability_cls = GenericDistributionMatchingCapability + @cached_property + def data_processor(self): + return QwenImageDataProcessor(self.image_processor, self.config) + def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -120,7 +126,7 @@ def _denormalize_latents(self, latents): return latents * std + mean def encode_to_latent(self, sample): - image = sample["inputs"]["target_image"].to(device=self.device, dtype=self.running_dtype) + image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) pixel_values = image.unsqueeze(2) latent = self.vae.encode(pixel_values).latent_dist.sample() # (B, C, T, H, W) return self._normalize_latents(latent) @@ -129,6 +135,9 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) + def infer_target_size(self, sample, default_height, default_width): + return self.data_processor.infer_target_size(sample, default_height, default_width) + def encode_prompt_condition(self, prompt, **kwargs): prompt_embed, prompt_embed_mask = self.text_pipeline.encode_prompt( prompt=prompt, diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index cfea1899e..4179fab52 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -1,4 +1,4 @@ -import math +from functools import cached_property import torch from PIL import Image @@ -7,19 +7,9 @@ from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageEditDistributionMatchingCapability from lightx2v_train.utils.registry import MODEL_REGISTER +from .data_process import QwenImageEditDataProcessor from .qwen_image import QwenImageModel -CONDITION_IMAGE_AREA = 384 * 384 -VAE_IMAGE_AREA = 1024 * 1024 - - -def _calculate_dimensions(target_area, ratio): - width = math.sqrt(target_area * ratio) - height = width / ratio - width = round(width / 32) * 32 - height = round(height / 32) * 32 - return int(width), int(height) - @MODEL_REGISTER("qwen_image_edit") class QwenImageEditModel(QwenImageModel): @@ -30,13 +20,16 @@ class QwenImageEditModel(QwenImageModel): pipeline_cls = QwenImageEditPlusPipeline distribution_matching_capability_cls = QwenImageEditDistributionMatchingCapability + @cached_property + def data_processor(self): + return QwenImageEditDataProcessor(self.image_processor, self.config) + def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_conditions_with_source(sample, [prompt])[0] def encode_conditions_with_source(self, sample, prompts): - source_images = self._source_images_from_sample(sample) - condition_images, vae_images = self._prepare_source_images(source_images) + condition_images, vae_images = self.data_processor.process_source_images(sample) conditions = [self.encode_prompt_condition(prompt, image=condition_images) for prompt in prompts] if vae_images: source_latents, source_img_shapes = self._encode_source_image_latents(vae_images) @@ -45,41 +38,6 @@ def encode_conditions_with_source(self, sample, prompts): condition["source_img_shapes"] = source_img_shapes return conditions - def _source_images_from_sample(self, sample): - source_images = sample["inputs"].get("source_images") - if source_images is None: - return [] - tensors = [] - for image in source_images: - if not isinstance(image, torch.Tensor): - raise TypeError(f"source_images must contain tensors after collation, got {type(image)}") - if image.ndim == 3: - image = image.unsqueeze(0) - if image.ndim != 4 or image.shape[0] != 1: - raise ValueError(f"QwenImageEditPlusPipeline requires source images with batch_size=1, got {tuple(image.shape)}") - tensors.append(image) - return tensors - - def _prepare_source_images(self, source_images): - if not source_images: - return None, [] - - condition_images = [] - vae_images = [] - for image in source_images: - pil_image = self._tensor_to_pil(image[0]) - ratio = pil_image.width / pil_image.height - condition_width, condition_height = _calculate_dimensions(CONDITION_IMAGE_AREA, ratio) - vae_width, vae_height = _calculate_dimensions(VAE_IMAGE_AREA, ratio) - condition_images.append(self.image_processor.resize(pil_image, condition_height, condition_width)) - vae_images.append(self.image_processor.preprocess(pil_image, vae_height, vae_width).unsqueeze(2)) - return condition_images, vae_images - - def _tensor_to_pil(self, image): - image = ((image.detach().float().cpu().clamp(-1, 1) + 1.0) * 127.5).round().byte() - array = image.permute(1, 2, 0).numpy() - return Image.fromarray(array, mode="RGB") - def _encode_source_image_latents(self, vae_images): packed_latents = [] img_shapes = [] diff --git a/lightx2v_train/lightx2v_train/utils/image_ops.py b/lightx2v_train/lightx2v_train/utils/image_ops.py new file mode 100644 index 000000000..071188605 --- /dev/null +++ b/lightx2v_train/lightx2v_train/utils/image_ops.py @@ -0,0 +1,47 @@ +import math + +import torch +from PIL import Image + +_BICUBIC = getattr(Image, "Resampling", Image).BICUBIC + + +def image_tensor_to_pil(image): + if not torch.is_tensor(image): + raise TypeError(f"Expected an image tensor, got {type(image)}") + if image.ndim == 4: + if image.shape[0] != 1: + raise ValueError(f"Image preprocessing requires batch_size=1, got {tuple(image.shape)}") + image = image[0] + if image.ndim != 3 or image.shape[0] != 3: + raise ValueError(f"Expected an RGB tensor with shape (3, H, W), got {tuple(image.shape)}") + if image.dtype != torch.uint8: + raise TypeError(f"Expected a raw uint8 image tensor, got {image.dtype}") + array = image.detach().cpu().permute(1, 2, 0).contiguous().numpy() + return Image.fromarray(array) + + +def calculate_area_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 = max(multiple, round(math.sqrt(target_area * ratio) / multiple) * multiple) + height = max(multiple, round(math.sqrt(target_area / ratio) / multiple) * multiple) + return width, height + + +def resize_and_center_crop(image, width, height): + scale = max(width / image.width, height / image.height) + scaled_width = max(width, round(image.width * scale)) + scaled_height = max(height, round(image.height * scale)) + image = image.resize((scaled_width, scaled_height), resample=_BICUBIC) + left = (scaled_width - width) // 2 + top = (scaled_height - height) // 2 + return image.crop((left, top, left + width, top + height)) + + +def align_dimension(value, multiple): + if value <= 0: + raise ValueError(f"Image dimensions must be positive, got {value}") + return max(multiple, value // multiple * multiple) From f67b59e754b1ca5decab7bc1c6004626eece6fd5 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 23 Aug 2026 16:43:53 +0000 Subject: [PATCH 18/34] update dataset --- lightx2v_train/infer.py | 9 +- .../lightx2v_train/data/__init__.py | 3 +- .../lightx2v_train/data/image_dataset.py | 24 ++-- lightx2v_train/lightx2v_train/infer/image.py | 9 +- .../flux2_dopsd_capability.py | 4 +- .../model_zoo/flux2/data_process.py | 35 +++++- .../model_zoo/flux2/flux2_dev.py | 16 +-- .../model_zoo/flux2/flux2_klein.py | 16 +-- .../model_zoo/longcat_image/data_process.py | 34 +++++- .../model_zoo/longcat_image/longcat_image.py | 16 +-- .../model_zoo/qwen_image/data_process.py | 115 ++++++++++-------- .../model_zoo/qwen_image/qwen_image.py | 16 +-- .../model_zoo/qwen_image/qwen_image_edit.py | 17 +-- .../lightx2v_train/trainers/dopsd.py | 19 ++- .../lightx2v_train/utils/image_ops.py | 8 ++ .../lightx2v_train/utils/registry.py | 26 +++- lightx2v_train/train.py | 7 +- 17 files changed, 236 insertions(+), 138 deletions(-) diff --git a/lightx2v_train/infer.py b/lightx2v_train/infer.py index d0e1c9ed2..25d561389 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 @@ -93,6 +93,7 @@ def main(): setup_logger(config) try: + sample_processor = build_sample_processor(config) model = build_model(config) model.load_components() _load_full_checkpoint_for_infer( @@ -106,7 +107,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..f7190c545 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): @@ -27,6 +27,7 @@ def __getattr__(name): __all__ = [ "build_data", + "build_sample_processor", "build_image_dataset", "build_libero_fastwam_dataset", "build_latent_dataset", diff --git a/lightx2v_train/lightx2v_train/data/image_dataset.py b/lightx2v_train/lightx2v_train/data/image_dataset.py index 7ec4a04d9..cf74a2047 100644 --- a/lightx2v_train/lightx2v_train/data/image_dataset.py +++ b/lightx2v_train/lightx2v_train/data/image_dataset.py @@ -2,8 +2,6 @@ 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 @@ -13,14 +11,16 @@ class ImageDataset(Dataset): - """Decode image records without applying model-specific transforms.""" + """Decode image records and run an injected sample processor.""" def __init__( self, metadata_paths, prompt_dropout_rate=0.0, + sample_processor=None, ): self.prompt_dropout_rate = prompt_dropout_rate + self.sample_processor = sample_processor self.samples = [] for path in metadata_paths: path = Path(path) @@ -36,7 +36,9 @@ def __getitem__(self, index): prompt = record["prompt"] if random.random() < self.prompt_dropout_rate: prompt = " " + return self.load_sample(record, prompt=prompt) + 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"]) @@ -52,7 +54,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": @@ -90,12 +97,14 @@ def _resolve_path(self, path, data_dir): @staticmethod def load_image(image_path): with Image.open(image_path) as image: - array = np.array(image.convert("RGB"), dtype=np.uint8, copy=True) - return torch.from_numpy(array).permute(2, 0, 1) + 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): + if sample_processor is None: + raise ValueError("image_dataset requires a sample_processor") + 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" @@ -106,6 +115,7 @@ def build_image_dataset(data_config_split, train_or_val="train"): dataset = ImageDataset( metadata_paths=[Path(p) for p in data_path], prompt_dropout_rate=prompt_dropout_rate, + sample_processor=sample_processor, ) 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 diff --git a/lightx2v_train/lightx2v_train/infer/image.py b/lightx2v_train/lightx2v_train/infer/image.py index 8d4f73e5a..25ff2b42f 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -28,7 +28,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() @@ -70,7 +75,7 @@ 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 = self.model.infer_target_size(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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py index a2a62eba9..3541c53e4 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py @@ -62,9 +62,9 @@ def prepare_training_batch( teacher_prompts: list[str], running_dtype: torch.dtype, ) -> DopsdPreparedBatch: - image = batch.get("inputs", {}).get("target_image") + image = batch.get("inputs", {}).get("target_pixel_values") if image is None: - raise ValueError("DOPSD requires inputs.target_image.") + raise ValueError("DOPSD requires inputs.target_pixel_values.") image = self._image(image, running_dtype) height, width = image.shape[-2:] initial_state, state_ids, latent_hw = self._initial_state(height, width, generator=None) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py index dea34d058..313766e8a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -1,11 +1,27 @@ 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, - image_tensor_to_pil, 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 = 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): @@ -45,13 +61,19 @@ def _explicit_target_size(sample): return height, width +@SAMPLE_PROCESSOR_REGISTER("flux2_dev") +@SAMPLE_PROCESSOR_REGISTER("flux2_klein") class Flux2DataProcessor: - def __init__(self, image_processor, config): - self.image_processor = image_processor + def __init__(self, config): + self.image_processor = Flux2ImageProcessor(vae_scale_factor=_size_multiple_from_config(config)) self.target_area = _target_area_from_config(config) - def process_target(self, sample): - image = image_tensor_to_pil(sample["inputs"]["target_image"]) + 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: @@ -61,7 +83,8 @@ def process_target(self, sample): width = align_dimension(width, multiple) height = align_dimension(height, multiple) image = resize_and_center_crop(image, width, height) - return self.image_processor.preprocess(image) + 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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index 8aac16bfe..7c1a36638 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from functools import cached_property import torch from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel @@ -12,7 +11,6 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel -from .data_process import Flux2DataProcessor @dataclass @@ -27,10 +25,6 @@ class Flux2DevDenoiserInput: class Flux2DevModel(BaseModel): pipeline_cls = Flux2Pipeline - @cached_property - def data_processor(self): - return Flux2DataProcessor(self.image_processor, self.config) - def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -121,7 +115,12 @@ def _denormalize_patch_latents(self, latents): return Flux2Pipeline._unpatchify_latents(latents) def encode_to_latent(self, sample): - image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) + 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 = self.vae.encode(image).latent_dist.sample() return self._normalize_patch_latents(latent) @@ -129,9 +128,6 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) - def infer_target_size(self, sample, default_height, default_width): - return self.data_processor.infer_target_size(sample, default_height, default_width) - def encode_prompt_condition(self, prompt): model_config = self.config["model"] prompt_embed, text_ids = self.text_pipeline.encode_prompt( diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 80a29c697..59be42911 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from functools import cached_property import torch from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel @@ -24,7 +23,6 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel -from .data_process import Flux2DataProcessor @dataclass @@ -39,10 +37,6 @@ class Flux2KleinDenoiserInput: class Flux2KleinModel(BaseModel): pipeline_cls = Flux2KleinPipeline - @cached_property - def data_processor(self): - return Flux2DataProcessor(self.image_processor, self.config) - def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -135,7 +129,12 @@ def _denormalize_patch_latents(self, latents): return Flux2KleinPipeline._unpatchify_latents(latents) def encode_to_latent(self, sample): - image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) + 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 = self.vae.encode(image).latent_dist.sample() return self._normalize_patch_latents(latent) @@ -143,9 +142,6 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) - def infer_target_size(self, sample, default_height, default_width): - return self.data_processor.infer_target_size(sample, default_height, default_width) - def encode_prompt_condition(self, prompt): model_config = self.config["model"] prompt_embed, text_ids = self.text_pipeline.encode_prompt( 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 index 49dbae0b1..341f6274a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py @@ -1,11 +1,27 @@ import torch +from diffusers import AutoencoderKL +from diffusers.image_processor import VaeImageProcessor from lightx2v_train.utils.image_ops import ( align_dimension, calculate_area_dimensions, - image_tensor_to_pil, 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): @@ -45,13 +61,18 @@ def _explicit_target_size(sample): return height, width +@SAMPLE_PROCESSOR_REGISTER("longcat_image") class LongCatImageDataProcessor: - def __init__(self, image_processor, config): - self.image_processor = image_processor + def __init__(self, config): + self.image_processor = VaeImageProcessor(vae_scale_factor=_size_multiple_from_config(config)) self.target_area = _target_area_from_config(config) - def process_target(self, sample): - image = image_tensor_to_pil(sample["inputs"]["target_image"]) + 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: @@ -61,7 +82,8 @@ def process_target(self, sample): width = align_dimension(width, multiple) height = align_dimension(height, multiple) image = resize_and_center_crop(image, width, height) - return self.image_processor.preprocess(image) + 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) 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 index 1102ad306..227299bc8 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from functools import cached_property import torch from diffusers import AutoencoderKL, LongCatImagePipeline @@ -14,7 +13,6 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel -from .data_process import LongCatImageDataProcessor @dataclass @@ -29,10 +27,6 @@ class LongCatImageDenoiserInput: class LongCatImageModel(BaseModel): pipeline_cls = LongCatImagePipeline - @cached_property - def data_processor(self): - return LongCatImageDataProcessor(self.image_processor, self.config) - def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -112,7 +106,12 @@ def vae_scale_factor(self): return 2 ** (len(self.vae.config.block_out_channels) - 1) def encode_to_latent(self, sample): - image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) + 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 = 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) @@ -122,9 +121,6 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) - def infer_target_size(self, sample, default_height, default_width): - return self.data_processor.infer_target_size(sample, default_height, default_width) - def encode_prompt_condition(self, prompt): rewrite_training = self.config.get( "enable_prompt_rewrite_training", diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py index 6ce883a19..f8a5c6e74 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py @@ -1,16 +1,33 @@ import torch +from diffusers import AutoencoderKLQwenImage +from diffusers.image_processor import VaeImageProcessor from lightx2v_train.utils.image_ops import ( align_dimension, calculate_area_dimensions, - image_tensor_to_pil, + pil_to_image_tensor, resize_and_center_crop, ) +from lightx2v_train.utils.registry import SAMPLE_PROCESSOR_REGISTER CONDITION_IMAGE_AREA = 384 * 384 VAE_IMAGE_AREA = 1024 * 1024 +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 = AutoencoderKLQwenImage.load_config(model_path, subfolder="vae") + value = 2 ** len(vae_config["temperal_downsample"]) * 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): preprocessing = config.get("model", {}).get("input_preprocessing", {}) if "target_area" in preprocessing: @@ -48,13 +65,18 @@ def _explicit_target_size(sample): return height, width +@SAMPLE_PROCESSOR_REGISTER("qwen_image") class QwenImageDataProcessor: - def __init__(self, image_processor, config): - self.image_processor = image_processor + def __init__(self, config): + self.image_processor = VaeImageProcessor(vae_scale_factor=_size_multiple_from_config(config)) self.target_area = _target_area_from_config(config) - def process_target(self, sample): - return self._process_target(sample) + 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 infer_target_size(self, sample, default_height, default_width): height, width = _explicit_target_size(sample) @@ -63,78 +85,73 @@ def infer_target_size(self, sample, default_height, default_width): multiple = int(self.image_processor.config.vae_scale_factor) return align_dimension(int(height), multiple), align_dimension(int(width), multiple) - def _process_target(self, sample, reference_tensor=None, area_multiple=None): - image = image_tensor_to_pil(sample["inputs"]["target_image"]) - reference = image if reference_tensor is None else image_tensor_to_pil(reference_tensor) + def _process_target(self, image, sample, reference_image=None, area_multiple=None): + reference = image if reference_image is None else reference_image width, height = self._resolve_target_size( sample, - reference, + reference.width / reference.height, self.target_area, area_multiple, ) image = resize_and_center_crop(image, width, height) - return self.image_processor.preprocess(image) + return self.image_processor.preprocess(image)[0] - def _resolve_target_size(self, sample, reference_image, target_area, area_multiple=None): + def _resolve_target_size(self, sample, ratio, target_area, area_multiple=None): height, width = _explicit_target_size(sample) multiple = int(self.image_processor.config.vae_scale_factor) if height is not None: return align_dimension(width, multiple), align_dimension(height, multiple) - ratio = reference_image.width / reference_image.height return calculate_area_dimensions(target_area, ratio, area_multiple or multiple) +@SAMPLE_PROCESSOR_REGISTER("qwen_image_edit") class QwenImageEditDataProcessor(QwenImageDataProcessor): - def process_target(self, sample): - source_images = sample["inputs"].get("source_images") - reference_tensor = source_images[-1] if source_images else None - return self._process_target( - sample, - reference_tensor=reference_tensor, - area_multiple=32, - ) + def __call__(self, sample): + inputs = sample["inputs"] + target_image = inputs.pop("target_image", None) + source_images = inputs.pop("source_images", []) + reference_image = source_images[-1] if source_images else None + + if reference_image is not None: + sample["meta"]["reference_image_height"] = reference_image.height + sample["meta"]["reference_image_width"] = reference_image.width + if target_image is not None: + inputs["target_pixel_values"] = self._process_target( + target_image, + sample, + reference_image=reference_image, + area_multiple=32, + ) + if source_images: + condition_images, vae_images = self._process_source_images(source_images) + inputs["source_condition_images"] = condition_images + inputs["source_vae_pixel_values"] = vae_images + return sample def infer_target_size(self, sample, default_height, default_width): - source_images = sample["inputs"].get("source_images") - if not source_images: + meta = sample.get("meta", {}) + reference_height = _optional_scalar(meta, "reference_image_height") + reference_width = _optional_scalar(meta, "reference_image_width") + if (reference_height is None) != (reference_width is None): + raise ValueError("meta.reference_image_height and meta.reference_image_width must be provided together") + if reference_height is None: return super().infer_target_size(sample, default_height, default_width) - reference_image = image_tensor_to_pil(source_images[-1]) width, height = self._resolve_target_size( sample, - reference_image, + reference_width / reference_height, int(default_height) * int(default_width), area_multiple=32, ) return height, width - def process_source_images(self, sample): - source_images = self._source_images(sample) - if not source_images: - return None, [] - + def _process_source_images(self, source_images): condition_images = [] vae_images = [] for image in source_images: - pil_image = image_tensor_to_pil(image) - ratio = pil_image.width / pil_image.height + ratio = image.width / image.height condition_width, condition_height = calculate_area_dimensions(CONDITION_IMAGE_AREA, ratio, multiple=32) vae_width, vae_height = calculate_area_dimensions(VAE_IMAGE_AREA, ratio, multiple=32) - condition_images.append(self.image_processor.resize(pil_image, condition_height, condition_width)) - vae_images.append(self.image_processor.preprocess(pil_image, vae_height, vae_width).unsqueeze(2)) + condition_image = self.image_processor.resize(image, condition_height, condition_width) + condition_images.append(pil_to_image_tensor(condition_image)) + vae_images.append(self.image_processor.preprocess(image, vae_height, vae_width)[0].unsqueeze(1)) return condition_images, vae_images - - @staticmethod - def _source_images(sample): - source_images = sample["inputs"].get("source_images") - if source_images is None: - return [] - tensors = [] - for image in source_images: - if not isinstance(image, torch.Tensor): - raise TypeError(f"source_images must contain tensors after collation, got {type(image)}") - if image.ndim == 3: - image = image.unsqueeze(0) - if image.ndim != 4 or image.shape[0] != 1: - raise ValueError(f"QwenImageEditPlusPipeline requires source images with batch_size=1, got {tuple(image.shape)}") - tensors.append(image) - return tensors diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 2adaf98eb..2a5e17844 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -1,5 +1,4 @@ from dataclasses import dataclass -from functools import cached_property import torch from diffusers import AutoencoderKLQwenImage, QwenImagePipeline, QwenImageTransformer2DModel @@ -12,7 +11,6 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from ..base import BaseModel -from .data_process import QwenImageDataProcessor @dataclass @@ -34,10 +32,6 @@ class QwenImageModel(BaseModel): pipeline_cls = QwenImagePipeline distribution_matching_capability_cls = GenericDistributionMatchingCapability - @cached_property - def data_processor(self): - return QwenImageDataProcessor(self.image_processor, self.config) - def register_capabilities(self): super().register_capabilities() self.capabilities.register( @@ -126,7 +120,12 @@ def _denormalize_latents(self, latents): return latents * std + mean def encode_to_latent(self, sample): - image = self.data_processor.process_target(sample).to(device=self.device, dtype=self.running_dtype) + 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) pixel_values = image.unsqueeze(2) latent = self.vae.encode(pixel_values).latent_dist.sample() # (B, C, T, H, W) return self._normalize_latents(latent) @@ -135,9 +134,6 @@ def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_prompt_condition(prompt) - def infer_target_size(self, sample, default_height, default_width): - return self.data_processor.infer_target_size(sample, default_height, default_width) - def encode_prompt_condition(self, prompt, **kwargs): prompt_embed, prompt_embed_mask = self.text_pipeline.encode_prompt( prompt=prompt, diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index 4179fab52..eb76b5900 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -1,13 +1,11 @@ -from functools import cached_property - import torch from PIL import Image from diffusers import QwenImageEditPlusPipeline from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageEditDistributionMatchingCapability +from lightx2v_train.utils.image_ops import image_tensor_to_pil from lightx2v_train.utils.registry import MODEL_REGISTER -from .data_process import QwenImageEditDataProcessor from .qwen_image import QwenImageModel @@ -20,16 +18,15 @@ class QwenImageEditModel(QwenImageModel): pipeline_cls = QwenImageEditPlusPipeline distribution_matching_capability_cls = QwenImageEditDistributionMatchingCapability - @cached_property - def data_processor(self): - return QwenImageEditDataProcessor(self.image_processor, self.config) - def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] return self.encode_conditions_with_source(sample, [prompt])[0] def encode_conditions_with_source(self, sample, prompts): - condition_images, vae_images = self.data_processor.process_source_images(sample) + inputs = sample["inputs"] + condition_tensors = inputs.get("source_condition_images", []) + condition_images = [image_tensor_to_pil(image) for image in condition_tensors] or None + vae_images = inputs.get("source_vae_pixel_values", []) conditions = [self.encode_prompt_condition(prompt, image=condition_images) for prompt in prompts] if vae_images: source_latents, source_img_shapes = self._encode_source_image_latents(vae_images) @@ -42,6 +39,10 @@ def _encode_source_image_latents(self, vae_images): packed_latents = [] img_shapes = [] for image in vae_images: + if image.ndim == 4: + image = image.unsqueeze(0) + if image.ndim != 5 or image.shape[0] != 1: + raise ValueError(f"Expected one source VAE image with shape [1, C, T, H, W], got {tuple(image.shape)}") image = image.to(device=self.device, dtype=self.running_dtype) latent = self.vae.encode(image).latent_dist.mode() latent = self._normalize_latents(latent) diff --git a/lightx2v_train/lightx2v_train/trainers/dopsd.py b/lightx2v_train/lightx2v_train/trainers/dopsd.py index 61f85e48b..7520ef580 100644 --- a/lightx2v_train/lightx2v_train/trainers/dopsd.py +++ b/lightx2v_train/lightx2v_train/trainers/dopsd.py @@ -164,18 +164,17 @@ def _build_train_reference_map(self): target_image = record.get("target_image") prompt = record.get("prompt") if target_image is not None and prompt is not None: - self._train_reference_map[str(prompt).strip()] = target_image + self._train_reference_map[str(prompt).strip()] = record def _resolve_teacher_reference_image(self, record, dataset): - target_image = record.get("target_image") - if target_image is not None: - if isinstance(target_image, (str, os.PathLike)): - return dataset.load_image(target_image) - return target_image - - self._build_train_reference_map() - fallback_path = self._train_reference_map.get(str(record.get("prompt", "")).strip()) - return None if fallback_path is None else dataset.load_image(fallback_path) + reference_record = record if record.get("target_image") is not None else None + if reference_record is None: + self._build_train_reference_map() + reference_record = self._train_reference_map.get(str(record.get("prompt", "")).strip()) + if reference_record is None: + return None + sample = dataset.load_sample(reference_record) + return sample["inputs"].get("target_pixel_values") def compute_loss_on_sample(self, sample, collect_trajectory=False): return self.objective.compute_loss(sample, collect_trajectory=collect_trajectory) diff --git a/lightx2v_train/lightx2v_train/utils/image_ops.py b/lightx2v_train/lightx2v_train/utils/image_ops.py index 071188605..b83cb7947 100644 --- a/lightx2v_train/lightx2v_train/utils/image_ops.py +++ b/lightx2v_train/lightx2v_train/utils/image_ops.py @@ -1,11 +1,19 @@ import math +import numpy as np import torch from PIL import Image _BICUBIC = getattr(Image, "Resampling", Image).BICUBIC +def pil_to_image_tensor(image): + if not isinstance(image, Image.Image): + raise TypeError(f"Expected a PIL image, got {type(image)}") + array = np.array(image.convert("RGB"), dtype=np.uint8, copy=True) + return torch.from_numpy(array).permute(2, 0, 1).contiguous() + + def image_tensor_to_pil(image): if not torch.is_tensor(image): raise TypeError(f"Expected an image tensor, got {type(image)}") diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index 7f6526388..c02615781 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -69,6 +69,7 @@ def merge(self, other_register): TRAINER_REGISTER = Register() INFERENCER_REGISTER = Register() DATA_REGISTER = Register() +SAMPLE_PROCESSOR_REGISTER = Register() _MODEL_MODULES = { @@ -115,6 +116,14 @@ def merge(self, other_register): "wan_ti2v_5b_ar_infer": "lightx2v_train.infer.video", } +_SAMPLE_PROCESSOR_MODULES = { + "flux2_dev": "lightx2v_train.model_zoo.flux2.data_process", + "flux2_klein": "lightx2v_train.model_zoo.flux2.data_process", + "longcat_image": "lightx2v_train.model_zoo.longcat_image.data_process", + "qwen_image": "lightx2v_train.model_zoo.qwen_image.data_process", + "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image.data_process", +} + def _ensure_registered(name, register, module_map): if name in register: @@ -162,7 +171,19 @@ def build_inferencer(config): return INFERENCER_REGISTER[name](config) -def build_data(config, train_or_val): +def build_sample_processor(config): + processor_config = config.get("data", {}).get("processor", {}) + name = processor_config.get("name", config["model"]["name"]) + _ensure_registered(name, SAMPLE_PROCESSOR_REGISTER, _SAMPLE_PROCESSOR_MODULES) + if name not in SAMPLE_PROCESSOR_REGISTER: + if "name" not in processor_config: + return None + available = ", ".join(sorted(SAMPLE_PROCESSOR_REGISTER.keys())) + raise ValueError(f"Unknown sample processor {name!r}. Available processors: {available}") + return SAMPLE_PROCESSOR_REGISTER[name](config) + + +def build_data(config, train_or_val, sample_processor=None): data_config = config.get("data", {}) if train_or_val not in data_config: available_splits = ", ".join(repr(k) for k in sorted(data_config.keys())) @@ -177,4 +198,5 @@ def build_data(config, train_or_val): if data_name not in DATA_REGISTER: available_names = ", ".join(sorted(DATA_REGISTER.keys())) raise ValueError(f"Unknown data {data_name!r}. Available data: {available_names}") - return DATA_REGISTER[data_name](data_config_split, train_or_val=train_or_val) + kwargs = {"sample_processor": sample_processor} if data_name == "image_dataset" else {} + return DATA_REGISTER[data_name](data_config_split, train_or_val=train_or_val, **kwargs) diff --git a/lightx2v_train/train.py b/lightx2v_train/train.py index e7d794e9a..ef036c521 100644 --- a/lightx2v_train/train.py +++ b/lightx2v_train/train.py @@ -3,7 +3,7 @@ import torch from loguru import logger -from lightx2v_train.data import build_data, prepare_data +from lightx2v_train.data import build_data, build_sample_processor, prepare_data 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 @@ -25,13 +25,14 @@ def main(): try: prepare_data(config) + sample_processor = build_sample_processor(config) model = build_model(config) model.load_components() - dataloader_train = build_data(config, train_or_val="train") + dataloader_train = build_data(config, train_or_val="train", sample_processor=sample_processor) dataloader_eval = None if config.get("inference", {}).get("infer_every_iters", None): - dataloader_eval = build_data(config, train_or_val="val") + dataloader_eval = build_data(config, train_or_val="val", sample_processor=sample_processor) trainer = build_trainer(config) trainer.set_model(model) From 67eec1f54ea4b6f3a5c3ab203c28a93da3333f27 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 23 Aug 2026 17:41:48 +0000 Subject: [PATCH 19/34] update longcat & remove cdm --- .../infer/longcat_image_edit_lora.yaml | 50 ++++ .../configs/train/dmd/flux2_dev_cdm_lora.yaml | 157 ------------ .../train/dmd/longcat_image_cdm_lora.yaml | 132 ---------- .../train/dmd/longcat_image_dmd_lora.yaml | 1 - .../train/dmd/qwen_image_cdm_lora.yaml | 144 ----------- .../train/flow/longcat_image_edit_lora.yaml | 87 +++++++ .../train/flow/longcat_image_lora.yaml | 8 +- .../configs/train/flow/longcat_image_t2i.yaml | 8 +- lightx2v_train/data_process/README.md | 4 +- lightx2v_train/lightx2v_train/infer/base.py | 12 +- lightx2v_train/lightx2v_train/infer/image.py | 6 +- .../lightx2v_train/infer/image_native.py | 17 +- .../model_capabilities/training.py | 1 - .../lightx2v_train/model_zoo/__init__.py | 2 + .../lightx2v_train/model_zoo/base.py | 7 + .../model_zoo/capability_adapters/common.py | 2 +- .../capability_adapters/__init__.py | 3 +- ...e_edit_distribution_matching_capability.py | 23 ++ .../model_zoo/longcat_image/data_process.py | 65 +++++ .../model_zoo/longcat_image/longcat_image.py | 129 ++++++---- .../longcat_image/longcat_image_edit.py | 91 +++++++ ...max_h3_distribution_matching_capability.py | 1 - .../lightx2v_train/trainers/dmd/runtime.py | 61 ----- .../lightx2v_train/trainers/dmd/trainer.py | 45 +--- .../trainers/phased_dmd/config.py | 2 - .../lightx2v_train/trainers/sgmd.py | 2 - .../lightx2v_train/tricks/__init__.py | 10 - lightx2v_train/lightx2v_train/tricks/cdm.py | 230 ------------------ .../lightx2v_train/utils/registry.py | 2 + lightx2v_train/tests/test_longcat_image.py | 78 ++++++ 30 files changed, 525 insertions(+), 855 deletions(-) create mode 100644 lightx2v_train/configs/infer/longcat_image_edit_lora.yaml delete mode 100644 lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml delete mode 100644 lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml delete mode 100644 lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml create mode 100644 lightx2v_train/configs/train/flow/longcat_image_edit_lora.yaml create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_edit_distribution_matching_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py delete mode 100644 lightx2v_train/lightx2v_train/tricks/cdm.py create mode 100644 lightx2v_train/tests/test_longcat_image.py 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/train/dmd/flux2_dev_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml deleted file mode 100644 index 7641917b9..000000000 --- a/lightx2v_train/configs/train/dmd/flux2_dev_cdm_lora.yaml +++ /dev/null @@ -1,157 +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: - # JSONL row: {"prompt": "...", "target_height": 512, "target_width": 512} - name: prompt_dataset - num_workers: 8 - bucket_by_size: true - prompt_dropout_rate: 0.0 - shuffle: true - data_path: - - /path/to/datasets/image_dmd/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: - - {value: [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/longcat_image_cdm_lora.yaml b/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml deleted file mode 100644 index cf3d969ab..000000000 --- a/lightx2v_train/configs/train/dmd/longcat_image_cdm_lora.yaml +++ /dev/null @@ -1,132 +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: - # JSONL row: {"prompt": "...", "target_height": 1024, "target_width": 1024} - name: prompt_dataset - num_workers: 8 - bucket_by_size: true - prompt_dropout_rate: 0.0 - shuffle: true - data_path: - - /path/to/datasets/image_dmd/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: - - {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 - 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 4fd658dbf..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,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 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 f268a6dfd..000000000 --- a/lightx2v_train/configs/train/dmd/qwen_image_cdm_lora.yaml +++ /dev/null @@ -1,144 +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: - # JSONL row: {"prompt": "...", "target_height": 1024, "target_width": 1024} - name: prompt_dataset - num_workers: 8 - bucket_by_size: true - prompt_dropout_rate: 0.0 - shuffle: true - data_path: - - /path/to/datasets/image_dmd/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: - - {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 - 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/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 3831edb0a..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,7 +36,7 @@ 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 @@ -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 dd6e0f24c..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,7 +50,7 @@ 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 @@ -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/data_process/README.md b/lightx2v_train/data_process/README.md index 200169bca..7394dae6c 100644 --- a/lightx2v_train/data_process/README.md +++ b/lightx2v_train/data_process/README.md @@ -12,7 +12,7 @@ data_process/ | 训练方式 | 数据集 | 输入 | | --- | --- | --- | -| 图像 DMD / CDM | `prompt_dataset` | prompt、目标高度和目标宽度 | +| 图像 DMD | `prompt_dataset` | prompt、目标高度和目标宽度 | | Flow | `video_dataset` | 视频路径和 caption 元数据，训练时在线编码 | | Flow | `latent_dataset` | 预先生成的视频、音频和文本缓存 | | DMD / AR-DMD | `prompt_dataset` | TXT prompt，训练时在线编码 | @@ -68,7 +68,7 @@ data: data_path: /path/to/prompts/train.txt ``` -图像 DMD / CDM 使用 JSON 或 JSONL。每条记录必须显式指定目标尺寸， +图像 DMD 使用 JSON 或 JSONL。每条记录必须显式指定目标尺寸， 不需要 `target_image`： ```json 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/image.py b/lightx2v_train/lightx2v_train/infer/image.py index 25ff2b42f..798232f8f 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -58,7 +58,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 @@ -78,7 +78,7 @@ def infer(self): 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 = { @@ -86,7 +86,7 @@ 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: diff --git a/lightx2v_train/lightx2v_train/infer/image_native.py b/lightx2v_train/lightx2v_train/infer/image_native.py index ad796921c..0dc8d7bb0 100644 --- a/lightx2v_train/lightx2v_train/infer/image_native.py +++ b/lightx2v_train/lightx2v_train/infer/image_native.py @@ -29,8 +29,14 @@ def infer(self): 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) + should_load_lora = lora_path and getattr(self.model, "_infer_lora_adapter_name", None) is None + lora_loaded_by_pipeline = False + if should_load_lora: + if hasattr(pipe, "load_lora_weights"): + pipe.load_lora_weights(lora_path) + lora_loaded_by_pipeline = True + else: + self.model.load_lora_for_infer(lora_path) saved_paths = [] self.model.set_denoiser_eval() @@ -52,7 +58,10 @@ def infer(self): logger.info("Saved to {}", save_path) saved_paths.append(str(save_path)) - if lora_path: - pipe.unload_lora_weights() + if should_load_lora: + if lora_loaded_by_pipeline: + pipe.unload_lora_weights() + else: + self.model.unload_lora_for_infer() return saved_paths diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 849ee6686..25d0b04be 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -101,7 +101,6 @@ class DistributionMatchingProfile: """Feature and precision constraints published by a distribution-matching adapter.""" supports_guidance: bool = True - supports_cdm: bool = True supports_ida: bool = True supports_diversity: bool = True supports_real_data_fake: bool = True diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index cec809811..594247965 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -7,6 +7,7 @@ "Flux2KleinModel": (".flux2.flux2_klein", "Flux2KleinModel"), "LingBotVideoModel": (".wan.lingbot_video", "LingBotVideoModel"), "LongCatImageModel": (".longcat_image.longcat_image", "LongCatImageModel"), + "LongCatImageEditModel": (".longcat_image.longcat_image_edit", "LongCatImageEditModel"), "LTX2T2AVModel": (".ltx.ltx_t2av", "LTX2T2AVModel"), "MiniMaxH3T2AVModel": (".minimax_h3.minimax_h3_t2av", "MiniMaxH3T2AVModel"), "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), @@ -49,6 +50,7 @@ def __getattr__(name): "QwenImageModel", "QwenImageEditModel", "LongCatImageModel", + "LongCatImageEditModel", "Flux2DevModel", "Flux2KleinModel", "LingBotVideoModel", diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 730627d4f..3598974c2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -119,6 +119,10 @@ def encode_to_latent(self, sample): def encode_condition(self, sample): raise NotImplementedError + 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 @@ -128,6 +132,9 @@ 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. diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py index 90f93c800..efa2b4211 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -372,7 +372,7 @@ def predict_guided_velocity( return self.predict_velocity(latents, sigma, condition) if self._guidance_in_denoiser_space: - denoiser_input = self.model.prepare_denoiser_input(latents) + denoiser_input = self.model.prepare_denoiser_input(latents, condition=condition) positive = self.model.denoise( denoiser_input, sigma, 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 index 585c6d6c5..a62dbc0aa 100644 --- 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 @@ -1,5 +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"] +__all__ = ["LongCatImageConsistencyModelCapability", "LongCatImageEditDistributionMatchingCapability"] 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..d6088ad45 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/capability_adapters/longcat_image_edit_distribution_matching_capability.py @@ -0,0 +1,23 @@ +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _negative_prompt, + _require_single_prompt, +) + + +class LongCatImageEditDistributionMatchingCapability(GenericDistributionMatchingCapability): + def encode_conditions(self, batch, negative_prompt, guidance_scale, broadcast): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + 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 index 341f6274a..2133da287 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py @@ -1,3 +1,5 @@ +import math + import torch from diffusers import AutoencoderKL from diffusers.image_processor import VaeImageProcessor @@ -5,6 +7,7 @@ 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 @@ -25,6 +28,9 @@ def _size_multiple_from_config(config): 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"]) @@ -61,6 +67,22 @@ def _explicit_target_size(sample): 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): @@ -91,3 +113,46 @@ def infer_target_size(self, sample, default_height, default_width): 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 __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 index 227299bc8..387aebec0 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -19,6 +19,7 @@ class LongCatImageDenoiserInput: hidden_states: torch.Tensor img_ids: torch.Tensor + target_token_length: int height: int width: int @@ -26,6 +27,9 @@ class LongCatImageDenoiserInput: @MODEL_REGISTER("longcat_image") class LongCatImageModel(BaseModel): pipeline_cls = LongCatImagePipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability + supports_cfg_renorm = True + supports_prompt_rewrite = True def register_capabilities(self): super().register_capabilities() @@ -35,7 +39,7 @@ def register_capabilities(self): ) self.capabilities.register( DistributionMatchingCapability, - GenericDistributionMatchingCapability( + self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( channels_path="vae.config.latent_channels", @@ -59,13 +63,17 @@ def load_components(self, transformer_only=False, reference_model=None): return model_path = self.config["model"]["pretrained_model_name_or_path"] - self.text_pipeline = LongCatImagePipeline.from_pretrained( + self.text_pipeline = self.pipeline_cls.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.vae = AutoencoderKL.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) @@ -74,10 +82,14 @@ def load_components(self, transformer_only=False, reference_model=None): 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) + 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", None) + attention_backend = self.config["model"].get("attention_backend") if attention_backend is not None: self.transformer.set_attention_backend(attention_backend) @@ -105,6 +117,16 @@ def fsdp2_shard_plan(self, fsdp_config): def vae_scale_factor(self): return 2 ** (len(self.vae.config.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): image = sample["inputs"]["target_pixel_values"] if image.ndim == 3: @@ -112,48 +134,68 @@ def encode_to_latent(self, sample): 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 = 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 + return self._normalize_latents(self.vae.encode(image).latent_dist.sample()) 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): - 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 _get_additional_image_condition(self, condition, position_start): + del condition, position_start + return None, None + 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( + 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=(self.text_pipeline.tokenizer_max_length, self.text_pipeline.tokenizer_max_length), - height=h // 2, - width=w // 2, + 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=packed, + hidden_states=hidden_states, img_ids=img_ids, - height=h, - width=w, + target_token_length=target_tokens.shape[1], + height=height, + width=width, ) def denoise(self, denoiser_input, timestep_or_sigma, condition): - return self.transformer( + prediction = self.transformer( hidden_states=denoiser_input.hidden_states, timestep=timestep_or_sigma, guidance=None, @@ -162,35 +204,40 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): 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 LongCatImagePipeline._unpack_latents( + 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_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) + latent_height = 2 * (int(height) // (self.vae_scale_factor * 2)) + latent_width = 2 * (int(width) // (self.vae_scale_factor * 2)) + shape = (1, self.vae.config.latent_channels, latent_height, latent_width) return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) 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) + image = self.vae.decode(self._denormalize_latents(latent)).sample return self.image_processor.postprocess(image, output_type="pil") def assemble_pipeline(self, scheduler=None): - return LongCatImagePipeline( + return self.pipeline_cls( tokenizer=self.text_pipeline.tokenizer, text_encoder=self.text_pipeline.text_encoder, text_processor=self.text_pipeline.text_processor, @@ -205,7 +252,7 @@ def get_pipeline_infer_kwargs(self, infer_config): "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, + "guidance_scale": infer_config.get("cfg_guidance_scale", 4.5) 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/longcat_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py new file mode 100644 index 000000000..ec8370d39 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py @@ -0,0 +1,91 @@ +from PIL import Image +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): + pipeline_cls = LongCatImageEditPipeline + distribution_matching_capability_cls = LongCatImageEditDistributionMatchingCapability + supports_cfg_renorm = False + supports_prompt_rewrite = False + + def encode_condition(self, sample): + prompt = sample["conditioning"]["prompt"] + return self.encode_conditions_with_source(sample, [prompt])[0] + + 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 + + def get_pipeline_infer_kwargs(self, infer_config): + enable_cfg = infer_config.get("enable_cfg", False) + return { + "num_inference_steps": infer_config.get("num_inference_steps", 50), + "guidance_scale": infer_config.get("cfg_guidance_scale", 4.5) if enable_cfg else 1.0, + } + + def get_pipeline_sample_kwargs(self, sample): + source_paths = sample["meta"].get("source_image_paths", []) + if len(source_paths) != 1: + raise ValueError(f"LongCat Image Edit requires exactly one source image, got {len(source_paths)}") + with Image.open(source_paths[0]) as image: + return {"image": image.convert("RGB").copy()} 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 index f8382017f..7440d8b5d 100644 --- 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 @@ -83,7 +83,6 @@ class MiniMaxH3DistributionMatchingCapability(GenericDistributionMatchingCapabil ) _PROFILE = DistributionMatchingProfile( supports_guidance=False, - supports_cdm=False, supports_ida=False, supports_diversity=False, supports_real_data_fake=False, diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py index 8c25a8312..711a33537 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py @@ -18,9 +18,6 @@ ) from lightx2v_train.schedulers import DMDFlowMatchingScheduler from lightx2v_train.tricks import ( - CdmStepContext, - CdmTrainerConstraints, - CdmTrick, IdaModelPair, IdaSetupContext, IdaStepContext, @@ -44,7 +41,6 @@ class _DmdRuntime(BaseTrainer): DistributionMatchingCapability, ) default_negative_prompt = None - supports_cdm = True supports_ida = True defer_ida_setup = False @@ -77,13 +73,6 @@ def __init__(self, config): self.fake_optimizer_adam_epsilon = self.fake_optimizer_config.get("adam_epsilon", self.optimizer_adam_epsilon) self.dmd_config = parsed.dmd - self.cdm_trick = CdmTrick.from_mapping(self.dmd_config.get("cdm", {})) - self.cdm_trick.validate_trainer( - CdmTrainerConstraints( - supported=self.supports_cdm, - ), - self.trainer_name, - ) self.ida_trick = ImplicitDistributionAlignmentTrick.from_mapping(self.dmd_config.get("ida", {})) if self.ida_trick.enabled and not self.supports_ida: raise ValueError(f"{self.trainer_name} does not support training.dmd.ida.") @@ -123,7 +112,6 @@ def set_model(self, model): def _validate_distribution_matching_profile(self, profile): model_name = self.model_config.get("name", type(self.model).__name__) requested_features = ( - (self.cdm_trick.enabled, profile.supports_cdm, "CDM"), (self.ida_trick.enabled, profile.supports_ida, "IDA"), (getattr(getattr(self, "diversity_trick", None), "enabled", False), profile.supports_diversity, "diversity loss"), (getattr(getattr(self, "real_data_fake_trick", None), "enabled", False), profile.supports_real_data_fake, "real-data fake loss"), @@ -136,10 +124,6 @@ def _validate_distribution_matching_profile(self, profile): if getattr(self, "warp_denoising_step", False) and not profile.supports_warped_denoising_schedule: raise ValueError(f"model={model_name!r} requires an unwarped base denoising schedule; set training.dmd.warp_denoising_step=false.") - @property - def cdm_enabled(self): - return self.cdm_trick.enabled - def _get_optimizer_config(self): return self.training_config["student"]["optimizer"] @@ -266,12 +250,6 @@ def setup(self, resume_ckpt_path=None): self.random_schedule_sigma_max, self.random_schedule_sampling_method, ) - if self.cdm_enabled: - logger.info( - "[train] dmd CDM enabled: weight={} warmup_iters={}", - self.cdm_trick.config.weight, - self.cdm_trick.config.warmup_iters, - ) @staticmethod def _count_trainable(module): @@ -387,45 +365,6 @@ def _sample_synced_int(self, low, high): dist.broadcast(value, src=0) return int(value.item()) - def _cdm_step_context( - self, - xt, - vt, - end_step_idx, - condition, - current_iter, - ): - return CdmStepContext( - trajectory_latent=xt, - trajectory_velocity=vt, - end_step_index=end_step_idx, - condition=condition, - current_iteration=current_iter, - device=self.student.device, - dtype=self.latent_dtype, - scheduler=self.scheduler, - predict_student_velocity=lambda latent, sigma, cond: self._predict_velocity( - self.student, - latent, - sigma, - cond, - ), - predict_fake_velocity=lambda latent, sigma, cond: self._predict_velocity( - self.fake, - latent, - sigma, - cond, - ), - predict_teacher_velocity=lambda latent, sigma, cond: self._predict_teacher_velocity( - latent, - sigma, - cond, - None, - ), - prepare_fake_model=lambda: self.fake.set_training(False), - dmd_loss=self._dmd_loss, - ) - def _iter_train_samples(self): epoch = 0 while True: diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py index 7ea364b8e..30af6db9d 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py @@ -307,11 +307,8 @@ def train(self): samples, stage="student", grad_accum_iters=grad_accum_iters, - current_iter=current_iter, ) loss_dmd_value = student_result["dmd"] - loss_cdm_value = student_result.get("cdm", 0.0) - cdm_weight = student_result.get("cdm_weight", 0.0) loss_div_value = student_result["div_loss"] loss_real_dmd_value = student_result["real_dmd"] loss_fake_value = 0.0 @@ -321,7 +318,6 @@ def train(self): samples, stage="fake", grad_accum_iters=grad_accum_iters, - current_iter=current_iter, ) loss_fake_value += fake_result["loss"] loss_fake_real_value += fake_result["fake_real"] @@ -331,7 +327,6 @@ def train(self): current_iter += 1 display_fake = reduce_mean(loss_fake_value) display_dmd = reduce_mean(loss_dmd_value) if loss_dmd_value is not None else None - display_cdm = reduce_mean(loss_cdm_value) display_div = reduce_mean(loss_div_value) display_real_dmd = reduce_mean(loss_real_dmd_value) display_fake_real = reduce_mean(loss_fake_real_value) @@ -358,9 +353,6 @@ def train(self): } if display_dmd is not None: metrics["train/dmd"] = display_dmd - if self.cdm_enabled: - metrics["train/cdm"] = display_cdm - metrics["train/cdm_weight"] = cdm_weight self.log_metrics(metrics, step=current_iter) if save_every_iters and current_iter % save_every_iters == 0: @@ -376,7 +368,6 @@ def _train_one_stage( samples, stage, grad_accum_iters, - current_iter=None, ): if stage == "student": optimizer = self.optimizer @@ -414,7 +405,6 @@ def _train_one_stage( conditions, stage=stage, initial_noise=initial_noise, - current_iter=current_iter, ) if isinstance(result, dict): loss = result["loss"] @@ -617,7 +607,6 @@ def forward_loss( conditions, stage, initial_noise=None, - current_iter=None, ): condition, negative_condition = conditions ( @@ -699,29 +688,10 @@ def forward_loss( x_pred_fake, x_pred_teacher, ) - result = { + return { "loss": loss_dmd, "dmd": loss_dmd.detach(), } - if not self.cdm_enabled: - return result - - end_sample, end_velocity, end_step_index = self._last_rollout_end - cdm_result = self.cdm_trick.student_loss( - self._cdm_step_context( - end_sample, - end_velocity, - end_step_index, - condition, - current_iter, - ) - ) - result.update( - loss=loss_dmd + cdm_result.loss, - cdm=cdm_result.metrics["cdm"], - cdm_weight=cdm_result.metrics["cdm_weight"], - ) - return result def sample_end_step(self): return self._sample_synced_int( @@ -753,8 +723,6 @@ def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): ) ) x0 = None - end_sample = None - end_velocity = None # Zoe runs every DMD role in eval mode. Gradients are controlled by # autograd contexts, not by module.train(), which keeps rollout # dropout and other training-only behavior deterministic. @@ -773,23 +741,12 @@ def run_back_simulation(self, condition, latent_shape, grad_enabled, xt=None): sigma, condition, ) - if grad_enabled and idx == end_step_idx: - end_sample = self.student.detach(xt) - end_velocity = self.student.detach(velocity) xt, x0 = self.student.step( self.scheduler, velocity, idx, xt, ) - - if grad_enabled: - self._last_rollout_end = ( - end_sample, - end_velocity, - end_step_idx, - ) - return self.student.to_dtype( x0, self.latent_dtype, diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py index 81bf62d43..957dca42e 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/config.py @@ -65,8 +65,6 @@ def from_mapping( max_train_iters, default_lora_target_modules, ): - if dmd_config.get("cdm", {}).get("enabled", False): - raise ValueError("phased_dmd does not support training.dmd.cdm.") phased = dmd_config.get("phased", {}) if not isinstance(phased, dict): raise ValueError("training.dmd.phased must be a mapping.") diff --git a/lightx2v_train/lightx2v_train/trainers/sgmd.py b/lightx2v_train/lightx2v_train/trainers/sgmd.py index c1ee2b1a7..a6e18cd31 100644 --- a/lightx2v_train/lightx2v_train/trainers/sgmd.py +++ b/lightx2v_train/lightx2v_train/trainers/sgmd.py @@ -69,8 +69,6 @@ def __init__(self, config): super().__init__(config) if self.fake_update_ratio != 1: raise ValueError("SGMD uses one student and one fake-score update per iteration; set training.dmd.fake_update_ratio=1.") - if self.cdm_enabled: - raise ValueError("SGMD does not support training.dmd.cdm.") def sample_end_step(self): end_step_idx = self._sample_synced_int( diff --git a/lightx2v_train/lightx2v_train/tricks/__init__.py b/lightx2v_train/lightx2v_train/tricks/__init__.py index cebc6cca4..18bed765f 100644 --- a/lightx2v_train/lightx2v_train/tricks/__init__.py +++ b/lightx2v_train/lightx2v_train/tricks/__init__.py @@ -1,10 +1,4 @@ from .base import TrainerTrick, TrickLossResult -from .cdm import ( - CdmConfig, - CdmStepContext, - CdmTrainerConstraints, - CdmTrick, -) from .diversity import ( DiversityConfig, DiversitySetupContext, @@ -28,10 +22,6 @@ ) __all__ = [ - "CdmConfig", - "CdmStepContext", - "CdmTrainerConstraints", - "CdmTrick", "DiversityConfig", "DiversitySetupContext", "DiversityStepContext", diff --git a/lightx2v_train/lightx2v_train/tricks/cdm.py b/lightx2v_train/lightx2v_train/tricks/cdm.py deleted file mode 100644 index 672495523..000000000 --- a/lightx2v_train/lightx2v_train/tricks/cdm.py +++ /dev/null @@ -1,230 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Callable, Mapping, Protocol - -import torch - -from .base import TrainerTrick, TrickLossResult - - -class CdmScheduler(Protocol): - """Scheduler operations required by consistency distribution matching.""" - - def sigma_at( - self, - step_index: int, - *, - device: torch.device, - dtype: torch.dtype, - ) -> torch.Tensor: ... - - def sample_renoise_sigma( - self, - *, - device: torch.device, - dtype: torch.dtype, - ) -> torch.Tensor: ... - - def add_noise( - self, - latent: torch.Tensor, - noise: torch.Tensor, - sigma: torch.Tensor, - ) -> torch.Tensor: ... - - def _expand_to_ndim( - self, - value: torch.Tensor, - ndim: int, - ) -> torch.Tensor: ... - - -VelocityPrediction = Callable[ - [torch.Tensor, torch.Tensor, Any], - torch.Tensor, -] -DmdLoss = Callable[..., torch.Tensor] - - -@dataclass(frozen=True) -class CdmConfig: - enabled: bool = False - weight: float = 1.0 - warmup_iters: int = 0 - norm_clip_min: float = 0.1 - - @classmethod - def from_mapping(cls, config: Mapping[str, Any]) -> "CdmConfig": - if not isinstance(config, Mapping): - raise ValueError("training.dmd.cdm must be a mapping.") - return cls( - enabled=bool(config.get("enabled", False)), - weight=float(config.get("weight", 1.0)), - warmup_iters=int(config.get("warmup_iters", 0)), - norm_clip_min=float(config.get("norm_clip_min", 0.1)), - ) - - -@dataclass(frozen=True) -class CdmTrainerConstraints: - supported: bool - - -@dataclass(frozen=True) -class CdmStepContext: - trajectory_latent: torch.Tensor - trajectory_velocity: torch.Tensor - end_step_index: int - condition: Any - current_iteration: int | None - device: torch.device - dtype: torch.dtype - scheduler: CdmScheduler - predict_student_velocity: VelocityPrediction - predict_fake_velocity: VelocityPrediction - predict_teacher_velocity: VelocityPrediction - prepare_fake_model: Callable[[], None] - dmd_loss: DmdLoss - - -class CdmTrick( - TrainerTrick[ - None, - CdmStepContext, - None, - None, - ] -): - """Consistency distribution matching for image DMD trainers.""" - - name = "consistency_distribution_matching" - - def __init__(self, config: CdmConfig): - super().__init__(enabled=config.enabled) - self.config = config - self.validate() - - @classmethod - def from_mapping( - cls, - config: Mapping[str, Any], - ) -> "CdmTrick": - return cls(CdmConfig.from_mapping(config)) - - def validate(self) -> None: - if self.config.weight < 0: - raise ValueError("training.dmd.cdm.weight must be non-negative.") - if self.config.warmup_iters < 0: - raise ValueError("training.dmd.cdm.warmup_iters must be non-negative.") - if self.config.norm_clip_min <= 0: - raise ValueError("training.dmd.cdm.norm_clip_min must be positive.") - - def validate_trainer( - self, - constraints: CdmTrainerConstraints, - trainer_name: str, - ) -> None: - if not self.enabled: - return - if not constraints.supported: - raise ValueError(f"{trainer_name} does not support training.dmd.cdm.") - - def effective_weight(self, current_iteration: int | None) -> float: - if self.config.warmup_iters <= 0 or current_iteration is None: - return self.config.weight - progress = min( - 1.0, - max( - 0.0, - float(current_iteration) / float(self.config.warmup_iters), - ), - ) - return progress * self.config.weight - - def student_loss( - self, - context: CdmStepContext, - ) -> TrickLossResult: - weight = self.effective_weight(context.current_iteration) - if not self.enabled or weight == 0: - zero = context.trajectory_latent.new_zeros( - (), - dtype=torch.float32, - ) - return TrickLossResult( - loss=zero, - metrics={ - "cdm": zero.detach(), - "cdm_weight": weight, - }, - ) - - trajectory_sigma = context.scheduler.sigma_at( - context.end_step_index, - device=context.device, - dtype=torch.float32, - ) - student_sigma = context.scheduler.sample_renoise_sigma( - device=context.device, - dtype=torch.float32, - ) - trajectory_sigma = context.scheduler._expand_to_ndim( - trajectory_sigma, - context.trajectory_velocity.ndim, - ) - student_sigma_expanded = context.scheduler._expand_to_ndim( - student_sigma, - context.trajectory_velocity.ndim, - ) - student_xt = context.trajectory_latent + (student_sigma_expanded - trajectory_sigma) * context.trajectory_velocity - student_prediction = context.predict_student_velocity( - student_xt.to(context.dtype), - student_sigma, - context.condition, - ) - student_x0 = (student_xt - student_sigma_expanded * student_prediction).to(context.dtype) - - teacher_sigma = context.scheduler.sample_renoise_sigma( - device=context.device, - dtype=context.dtype, - ) - teacher_noise = torch.randn_like(student_x0) - teacher_xt = context.scheduler.add_noise( - student_x0, - teacher_noise, - teacher_sigma, - ) - with torch.no_grad(): - context.prepare_fake_model() - velocity_fake = context.predict_fake_velocity( - teacher_xt, - teacher_sigma, - context.condition, - ) - velocity_teacher = context.predict_teacher_velocity( - teacher_xt, - teacher_sigma, - context.condition, - ) - - teacher_sigma_expanded = context.scheduler._expand_to_ndim( - teacher_sigma, - teacher_xt.ndim, - ) - x_pred_fake = teacher_xt - teacher_sigma_expanded * velocity_fake - x_pred_teacher = teacher_xt - teacher_sigma_expanded * velocity_teacher - - raw_loss = context.dmd_loss( - student_x0, - x_pred_fake, - x_pred_teacher, - norm_clip_min=self.config.norm_clip_min, - ) - return TrickLossResult( - loss=weight * raw_loss, - metrics={ - "cdm": raw_loss.detach(), - "cdm_weight": weight, - }, - ) diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index c02615781..4ee582a78 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -77,6 +77,7 @@ def merge(self, other_register): "flux2_klein": "lightx2v_train.model_zoo.flux2.flux2_klein", "lingbot_video": "lightx2v_train.model_zoo.wan.lingbot_video", "longcat_image": "lightx2v_train.model_zoo.longcat_image.longcat_image", + "longcat_image_edit": "lightx2v_train.model_zoo.longcat_image.longcat_image_edit", "ltx_t2av": "lightx2v_train.model_zoo.ltx.ltx_t2av", "ltx_t2av_ar": "lightx2v_train.model_zoo.ltx.ltx_t2av", "minimax_h3_t2av": "lightx2v_train.model_zoo.minimax_h3.minimax_h3_t2av", @@ -120,6 +121,7 @@ def merge(self, other_register): "flux2_dev": "lightx2v_train.model_zoo.flux2.data_process", "flux2_klein": "lightx2v_train.model_zoo.flux2.data_process", "longcat_image": "lightx2v_train.model_zoo.longcat_image.data_process", + "longcat_image_edit": "lightx2v_train.model_zoo.longcat_image.data_process", "qwen_image": "lightx2v_train.model_zoo.qwen_image.data_process", "qwen_image_edit": "lightx2v_train.model_zoo.qwen_image.data_process", } diff --git a/lightx2v_train/tests/test_longcat_image.py b/lightx2v_train/tests/test_longcat_image.py new file mode 100644 index 000000000..4ff8a60ab --- /dev/null +++ b/lightx2v_train/tests/test_longcat_image.py @@ -0,0 +1,78 @@ +from types import SimpleNamespace + +import torch +from PIL import Image + +from lightx2v_train.model_zoo.longcat_image.data_process import LongCatImageEditDataProcessor +from lightx2v_train.model_zoo.longcat_image.longcat_image import LongCatImageModel +from lightx2v_train.model_zoo.longcat_image.longcat_image_edit import LongCatImageEditModel + + +def _config(name="longcat_image_edit"): + return { + "model": { + "name": name, + "pretrained_model_name_or_path": "/unused", + "running_dtype": "fp32", + }, + "data": {"processor": {"size_multiple": 16, "target_area": 1024 * 1024}}, + "inference": {"enable_cfg_renorm": True, "cfg_renorm_min": 0.0}, + } + + +def test_edit_processor_builds_aligned_source_conditions(): + processor = LongCatImageEditDataProcessor(_config()) + sample = { + "inputs": { + "target_image": Image.new("RGB", (200, 100)), + "source_images": [Image.new("RGB", (200, 100))], + }, + "conditioning": {"prompt": "edit"}, + "meta": {}, + } + + sample = processor(sample) + inputs = sample["inputs"] + height, width = sample["meta"]["target_height"], sample["meta"]["target_width"] + + assert height % 16 == width % 16 == 0 + assert inputs["target_pixel_values"].shape == (3, height, width) + assert inputs["source_vae_pixel_values"].shape == (3, height, width) + assert inputs["source_condition_image"].shape == (3, height // 2, width // 2) + assert inputs["source_condition_image"].dtype == torch.uint8 + assert processor.infer_target_size(sample, 1, 1) == (height, width) + + +def test_edit_denoiser_input_appends_source_modality(): + model = LongCatImageEditModel(_config()) + model.vae = SimpleNamespace(config=SimpleNamespace(block_out_channels=[1, 1, 1, 1])) + noisy_latent = torch.randn(1, 16, 8, 12) + condition = { + "prompt_embed": torch.zeros(1, 520, 8), + "text_ids": torch.zeros(520, 3), + "source_tokens": torch.randn(1, 24, 64), + "source_height": 8, + "source_width": 12, + } + + denoiser_input = model.prepare_denoiser_input(noisy_latent, condition) + + assert denoiser_input.hidden_states.shape == (1, 48, 64) + assert denoiser_input.target_token_length == 24 + assert torch.all(denoiser_input.img_ids[:24, 0] == 1) + assert torch.all(denoiser_input.img_ids[24:, 0] == 2) + assert torch.all(denoiser_input.img_ids[:, 1:] >= 520) + + +def test_cfg_renorm_is_only_used_for_text_to_image(): + positive = torch.tensor([[[3.0, 4.0]]]) + negative = torch.zeros_like(positive) + + text_to_image = LongCatImageModel(_config("longcat_image")) + image_edit = LongCatImageEditModel(_config()) + + renormed = text_to_image.apply_cfg(positive, negative, 2.0) + plain = image_edit.apply_cfg(positive, negative, 2.0) + + assert torch.allclose(torch.linalg.vector_norm(renormed, dim=-1), torch.tensor([[5.0]])) + assert torch.equal(plain, positive * 2) From fbf764be13057d46ca937da45ad6576a70d5d861 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 23 Aug 2026 18:26:18 +0000 Subject: [PATCH 20/34] update flux & remove native infer --- .../consistency/qwen_image_pcm_cd_lora.yaml | 2 +- .../train/dmd/flux2_klein_dmd_lora.yaml | 2 +- .../train/dopsd/flux2_klein_dopsd.yaml | 6 +- .../train/flow/flux2_dev_edit_lora.yaml | 106 ++++++++ .../configs/train/flow/flux2_dev_lora.yaml | 4 +- .../train/flow/flux2_klein_edit_lora.yaml | 91 +++++++ .../configs/train/flow/flux2_klein_lora.yaml | 13 +- .../configs/train/flow/flux2_klein_t2i.yaml | 5 +- .../lightx2v_train/infer/__init__.py | 2 - lightx2v_train/lightx2v_train/infer/image.py | 18 +- .../lightx2v_train/infer/image_native.py | 67 ------ .../lightx2v_train/model_zoo/__init__.py | 4 + .../lightx2v_train/model_zoo/base.py | 16 -- .../flux2/capability_adapters/__init__.py | 7 +- ...2_edit_distribution_matching_capability.py | 25 ++ .../lightx2v_train/model_zoo/flux2/common.py | 226 ++++++++++++++++++ .../model_zoo/flux2/data_process.py | 103 ++++++-- .../model_zoo/flux2/flux2_dev.py | 182 ++------------ .../model_zoo/flux2/flux2_dev_edit.py | 11 + .../model_zoo/flux2/flux2_klein.py | 204 ++-------------- .../model_zoo/flux2/flux2_klein_edit.py | 12 + .../model_zoo/longcat_image/longcat_image.py | 22 -- .../longcat_image/longcat_image_edit.py | 15 -- .../model_zoo/qwen_image/qwen_image.py | 35 --- .../model_zoo/qwen_image/qwen_image_edit.py | 23 -- .../model_zoo/wan/lingbot_video.py | 13 - .../lightx2v_train/model_zoo/wan/wan_t2v.py | 20 -- .../lightx2v_train/utils/registry.py | 5 +- lightx2v_train/tests/test_longcat_image.py | 78 ------ 29 files changed, 637 insertions(+), 680 deletions(-) create mode 100644 lightx2v_train/configs/train/flow/flux2_dev_edit_lora.yaml create mode 100644 lightx2v_train/configs/train/flow/flux2_klein_edit_lora.yaml delete mode 100644 lightx2v_train/lightx2v_train/infer/image_native.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_edit_distribution_matching_capability.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/common.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py delete mode 100644 lightx2v_train/tests/test_longcat_image.py 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 index 6846da869..1683e3a6e 100644 --- a/lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml +++ b/lightx2v_train/configs/train/consistency/qwen_image_pcm_cd_lora.yaml @@ -81,7 +81,7 @@ training: output_dir: ./output_train/qwen_image_pcm_cd_lora inference: - method: image_native_infer + method: image_infer negative_prompt: " " default_width: 1024 default_height: 1024 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 5e13abb71..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,6 +1,6 @@ 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 diff --git a/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml b/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml index cc0f88895..1e3ecd2b6 100644 --- a/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml +++ b/lightx2v_train/configs/train/dopsd/flux2_klein_dopsd.yaml @@ -1,22 +1,22 @@ 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 text_encoder_out_layers: [9, 18, 27] running_dtype: bf16 data: + processor: + target_area: 1048576 # 1024 * 1024 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 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 8d32cbd91..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 @@ -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 6375c2c12..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,7 +36,7 @@ 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] @@ -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 5ef54d074..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: 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/image.py b/lightx2v_train/lightx2v_train/infer/image.py index 798232f8f..b224f34c0 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -15,6 +15,18 @@ def _record_has_source_images(record): @INFERENCER_REGISTER("image_infer") class ImageInferencer(BaseInferencer): + def _inference_sigmas(self, num_inference_steps): + 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)) + def _load_infer_sample(self, index, prompt): infer_sample = self.dataloader_eval.dataset[index] infer_sample["conditioning"]["prompt"] = prompt @@ -93,7 +105,11 @@ def infer(self): 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, + ) 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 0dc8d7bb0..000000000 --- a/lightx2v_train/lightx2v_train/infer/image_native.py +++ /dev/null @@ -1,67 +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 - should_load_lora = lora_path and getattr(self.model, "_infer_lora_adapter_name", None) is None - lora_loaded_by_pipeline = False - if should_load_lora: - if hasattr(pipe, "load_lora_weights"): - pipe.load_lora_weights(lora_path) - lora_loaded_by_pipeline = True - else: - self.model.load_lora_for_infer(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 should_load_lora: - if lora_loaded_by_pipeline: - pipe.unload_lora_weights() - else: - self.model.unload_lora_for_infer() - - return saved_paths diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 594247965..8e1b7eace 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -4,7 +4,9 @@ _LAZY_EXPORTS = { "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"), @@ -52,7 +54,9 @@ def __getattr__(name): "LongCatImageModel", "LongCatImageEditModel", "Flux2DevModel", + "Flux2DevEditModel", "Flux2KleinModel", + "Flux2KleinEditModel", "LingBotVideoModel", "LTX2T2AVModel", "MiniMaxH3T2AVModel", diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 3598974c2..834bb2703 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -169,22 +169,6 @@ def prepare_infer_latents(self, height, width, generator=None): 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: 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 index e8442a789..9e44bed60 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/__init__.py @@ -2,5 +2,10 @@ from .flux2_consistency_model_capability import Flux2ConsistencyModelCapability from .flux2_dopsd_capability import Flux2DopsdCapability +from .flux2_edit_distribution_matching_capability import Flux2EditDistributionMatchingCapability -__all__ = ["Flux2ConsistencyModelCapability", "Flux2DopsdCapability"] +__all__ = [ + "Flux2ConsistencyModelCapability", + "Flux2DopsdCapability", + "Flux2EditDistributionMatchingCapability", +] 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..aa0702d86 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_edit_distribution_matching_capability.py @@ -0,0 +1,25 @@ +import torch + +from lightx2v_train.model_zoo.capability_adapters.common import ( + GenericDistributionMatchingCapability, + _negative_prompt, + _require_single_prompt, +) + + +class Flux2EditDistributionMatchingCapability(GenericDistributionMatchingCapability): + """Keep reference-image conditioning in Flux2 edit DMD.""" + + def encode_conditions(self, batch, negative_prompt, guidance_scale, broadcast): + conditioning = batch["conditioning"] + prompt = conditioning.get("prompt", "") + scalar = _require_single_prompt(prompt) + 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..b005fbfcd --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py @@ -0,0 +1,226 @@ +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 = () + + def load_components(self, transformer_only=False, reference_model=None): + model_path = self.config["model"]["pretrained_model_name_or_path"] + self._validate_model_path(model_path) + self._configure_model() + + 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._copy_model_state(reference_model) + self.transformer = self.load_transformer(model_path) + return + + self.text_pipeline = self.pipeline_cls.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(model_path) + + 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 _validate_model_path(self, model_path): + del model_path + + def _configure_model(self): + pass + + def _copy_model_state(self, reference_model): + del reference_model + + 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): + return 2 ** (len(self.vae.config.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): + 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, self.target_latent_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_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 index 313766e8a..2f9fc5653 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -2,18 +2,13 @@ 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.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", {}) - preprocessing = config.get("model", {}).get("input_preprocessing", {}) - value = processor_config.get("size_multiple", preprocessing.get("size_multiple")) + 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") @@ -25,14 +20,8 @@ def _size_multiple_from_config(config): def _target_area_from_config(config): - 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 + processor_config = config.get("data", {}).get("processor", {}) + return int(processor_config.get("target_area", 1024 * 1024)) def _optional_scalar(mapping, key): @@ -71,20 +60,26 @@ def __init__(self, config): def __call__(self, sample): inputs = sample["inputs"] image = inputs.pop("target_image", None) - if image is None: - return sample + 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 None: - ratio = image.width / image.height - width, height = calculate_area_dimensions(self.target_area, ratio, multiple) - else: - width = align_dimension(width, multiple) + 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) - inputs["target_pixel_values"] = self.image_processor.preprocess(image)[0] - return sample + return self.image_processor.preprocess(image)[0] def infer_target_size(self, sample, default_height, default_width): height, width = _explicit_target_size(sample) @@ -92,3 +87,63 @@ def infer_target_size(self, sample, default_height, default_width): 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"] + 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] + sample["meta"]["reference_image_height"] = reference_height + sample["meta"]["reference_image_width"] = reference_width + inputs["source_vae_pixel_values"] = source_pixels + if target_image is not None: + inputs["target_pixel_values"] = self._process_target( + target_image, + sample, + fallback_size=(reference_height, reference_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 index 7c1a36638..e3ffd92d9 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -1,8 +1,5 @@ -from dataclasses import dataclass - import torch -from diffusers import AutoencoderKLFlux2, Flux2Pipeline, Flux2Transformer2DModel -from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from diffusers import Flux2Pipeline from lightx2v_train.model_capabilities import ConsistencyModelCapability, DistributionMatchingCapability, FlowMatchingSFTCapability from lightx2v_train.model_zoo.capability_adapters import SpatialLatentGeometry @@ -10,30 +7,21 @@ from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyModelCapability 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 +from .common import Flux2ModelBase @MODEL_REGISTER("flux2_dev") -class Flux2DevModel(BaseModel): +class Flux2DevModel(Flux2ModelBase): 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(FlowMatchingSFTCapability, GenericFlowMatchingCapability(self)) self.capabilities.register( DistributionMatchingCapability, - GenericDistributionMatchingCapability( + self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( channels_path="transformer.config.in_channels", @@ -42,113 +30,13 @@ def register_capabilities(self): guidance_in_denoiser_space=True, ), ) - self.capabilities.register( - ConsistencyModelCapability, - Flux2ConsistencyModelCapability(self), - ) - - 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 + self.capabilities.register(ConsistencyModelCapability, Flux2ConsistencyModelCapability(self)) - 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"], - }, - ] + def _configure_model(self): + self.guidance_scale = float(self.config["model"].get("guidance_scale", 4.0)) - @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_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 = 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 _copy_model_state(self, reference_model): + self.guidance_scale = reference_model.guidance_scale def denoise(self, denoiser_input, timestep_or_sigma, condition): guidance = torch.full( @@ -157,46 +45,8 @@ def denoise(self, denoiser_input, timestep_or_sigma, condition): 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 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) + return self._denoise(denoiser_input, timestep_or_sigma, condition, guidance) - 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))), - } + 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..f68bd5def --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py @@ -0,0 +1,11 @@ +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): + requires_source_images = True + target_latent_mode = "mode" + distribution_matching_capability_cls = Flux2EditDistributionMatchingCapability diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 59be42911..0b16a29ea 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -1,8 +1,4 @@ -from dataclasses import dataclass - -import torch -from diffusers import AutoencoderKLFlux2, Flux2KleinPipeline, Flux2Transformer2DModel -from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor +from diffusers import Flux2KleinPipeline from lightx2v_train.model_capabilities import ( AdapterBankCapability, @@ -11,41 +7,29 @@ DopsdCapability, FlowMatchingSFTCapability, ) -from lightx2v_train.model_zoo.capability_adapters import ( - PeftAdapterBankCapability, - SpatialLatentGeometry, -) +from lightx2v_train.model_zoo.capability_adapters import PeftAdapterBankCapability, SpatialLatentGeometry from lightx2v_train.model_zoo.capability_adapters.common import GenericDistributionMatchingCapability, GenericFlowMatchingCapability -from lightx2v_train.model_zoo.flux2.capability_adapters import ( - Flux2ConsistencyModelCapability, - Flux2DopsdCapability, -) +from lightx2v_train.model_zoo.flux2.capability_adapters import Flux2ConsistencyModelCapability, Flux2DopsdCapability 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 +from .common import Flux2ModelBase @MODEL_REGISTER("flux2_klein") -class Flux2KleinModel(BaseModel): +class Flux2KleinModel(Flux2ModelBase): + """Flux2 Klein Base model. Step-distilled Klein checkpoints are unsupported.""" + pipeline_cls = Flux2KleinPipeline + distribution_matching_capability_cls = GenericDistributionMatchingCapability + default_text_encoder_out_layers = (9, 18, 27) + supports_dopsd = True def register_capabilities(self): super().register_capabilities() - self.capabilities.register( - FlowMatchingSFTCapability, - GenericFlowMatchingCapability(self), - ) + self.capabilities.register(FlowMatchingSFTCapability, GenericFlowMatchingCapability(self)) self.capabilities.register( DistributionMatchingCapability, - GenericDistributionMatchingCapability( + self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( channels_path="transformer.config.in_channels", @@ -53,165 +37,23 @@ def register_capabilities(self): ), ), ) - self.capabilities.register( - ConsistencyModelCapability, - Flux2ConsistencyModelCapability(self), - ) - self.capabilities.register( - DopsdCapability, - Flux2DopsdCapability(self), - ) - self.capabilities.register( - AdapterBankCapability, - PeftAdapterBankCapability(self), - ) - - 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) + self.capabilities.register(ConsistencyModelCapability, Flux2ConsistencyModelCapability(self)) + if self.supports_dopsd: + self.capabilities.register(DopsdCapability, Flux2DopsdCapability(self)) + self.capabilities.register(AdapterBankCapability, PeftAdapterBankCapability(self)) - 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_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 = 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 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 _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.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] + return self._denoise(denoiser_input, timestep_or_sigma, condition, guidance=None) def postprocess_denoiser_output(self, prediction, denoiser_input): - return Flux2KleinPipeline._unpack_latents_with_ids( + return self.pipeline_cls._unpack_latents_with_ids( prediction, - denoiser_input.img_ids, + denoiser_input.target_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") - - 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))), - } 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..aef8d2dd0 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py @@ -0,0 +1,12 @@ +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): + requires_source_images = True + target_latent_mode = "mode" + distribution_matching_capability_cls = Flux2EditDistributionMatchingCapability + supports_dopsd = False 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 index 387aebec0..1c8bcb7ae 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -235,25 +235,3 @@ def prepare_infer_latents(self, height, width, generator=None): def decode_latent(self, latent): image = self.vae.decode(self._denormalize_latents(latent)).sample return self.image_processor.postprocess(image, output_type="pil") - - def assemble_pipeline(self, scheduler=None): - return self.pipeline_cls( - 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.5) 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/longcat_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image_edit.py index ec8370d39..d90210cb1 100644 --- 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 @@ -1,4 +1,3 @@ -from PIL import Image from diffusers import LongCatImageEditPipeline from diffusers.pipelines.longcat_image.pipeline_longcat_image import prepare_pos_ids @@ -75,17 +74,3 @@ def _get_additional_image_condition(self, condition, position_start): width=int(condition["source_width"]) // 2, ).to(self.device) return source_tokens, source_ids - - def get_pipeline_infer_kwargs(self, infer_config): - enable_cfg = infer_config.get("enable_cfg", False) - return { - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", 4.5) if enable_cfg else 1.0, - } - - def get_pipeline_sample_kwargs(self, sample): - source_paths = sample["meta"].get("source_image_paths", []) - if len(source_paths) != 1: - raise ValueError(f"LongCat Image Edit requires exactly one source image, got {len(source_paths)}") - with Image.open(source_paths[0]) as image: - return {"image": image.convert("RGB").copy()} diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 2a5e17844..50bc4bd6d 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -197,38 +197,3 @@ def decode_latent(self, latent): image = self.vae.decode(latent).sample # (B, C, T, H, W) image = image[:, :, 0, :, :] # drop temporal dim -> (B, C, H, W), T == 1 return self.image_processor.postprocess(image, output_type="pil") - - def get_pipeline_extra_components(self): - return {} - - def assemble_pipeline(self, scheduler=None): - components = { - "tokenizer": self.text_pipeline.tokenizer, - "text_encoder": self.text_pipeline.text_encoder, - "vae": self.vae, - "transformer": self.transformer, - "scheduler": scheduler if scheduler is not None else self.text_pipeline.scheduler, - **self.get_pipeline_extra_components(), - } - return self.pipeline_cls(**components).to(self.device) - - def get_pipeline_infer_kwargs(self, infer_config): - # QwenImagePipeline uses `true_cfg_scale` instead of the standard `guidance_scale` - enable_cfg = infer_config.get("enable_cfg", True) - kwargs = { - "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), - "true_cfg_scale": infer_config.get("cfg_guidance_scale", 4.0) if enable_cfg else 1.0, - } - sigmas = infer_config.get("sigmas") - pcm_solver_steps = 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 not None: - from lightx2v_train.trainers.consistency.pcm import pcm_inference_sigmas - - sigmas = pcm_inference_sigmas(kwargs["num_inference_steps"], int(pcm_solver_steps)) - if sigmas is not None: - kwargs["sigmas"] = [float(value) for value in sigmas] - return kwargs diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index eb76b5900..81cc54d0a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -1,5 +1,4 @@ import torch -from PIL import Image from diffusers import QwenImageEditPlusPipeline from lightx2v_train.model_zoo.qwen_image.capability_adapters import QwenImageEditDistributionMatchingCapability @@ -57,25 +56,3 @@ def _get_additional_image_tokens(self, condition): if condition is None: raise ValueError("QwenImageEditModel.prepare_denoiser_input requires condition.") return condition.get("source_latents"), condition.get("source_img_shapes", []) - - def get_pipeline_extra_components(self): - return {"processor": self.text_pipeline.processor} - - def get_pipeline_infer_kwargs(self, infer_config): - kwargs = super().get_pipeline_infer_kwargs(infer_config) - if infer_config.get("height") is None: - kwargs.pop("height") - if infer_config.get("width") is None: - kwargs.pop("width") - return kwargs - - def get_pipeline_sample_kwargs(self, sample): - source_image_paths = sample["meta"].get("source_image_paths", []) - if not source_image_paths: - return {} - - images = [] - for path in source_image_paths: - with Image.open(path) as image: - images.append(image.convert("RGB").copy()) - return {"image": images[0] if len(images) == 1 else images} diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index 12a07382f..b7b942190 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -445,19 +445,6 @@ def decode_latent(self, latent): frames = frames.permute(0, 2, 3, 4, 1).cpu().numpy() return [video for video in frames] - def assemble_pipeline(self, scheduler=None): - raise NotImplementedError("LingBot-Video training currently supports DMD only; native training inference is disabled.") - - def get_pipeline_infer_kwargs(self, infer_config): - return { - "height": infer_config.get("default_height", 480), - "width": infer_config.get("default_width", 832), - "num_frames": infer_config.get("num_frames", 81), - "num_inference_steps": infer_config.get("num_inference_steps", 40), - "guidance_scale": infer_config.get("cfg_guidance_scale", 3.0), - "shift": infer_config.get("shift", 3.0), - } - def log_model_structure(self): logger.info("[model] class={}", self.__class__.__name__) logger.info("[model] LingBot-Video transformer blocks={} experts_per_block={}", len(self.transformer.blocks), self.transformer.config.num_experts) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py index ae1171e1e..bd16f02a2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py @@ -453,26 +453,6 @@ def decode_latent(self, latent): videos = torch.stack(self.vae.decode(self._batch_to_list(latent)), dim=0) return self._postprocess_videos(videos) - def assemble_pipeline(self, scheduler=None): - raise NotImplementedError("Native Wan T2V uses wan_t2v_infer instead of assembling a diffusers pipeline.") - - def get_pipeline_infer_kwargs(self, infer_config): - if "height" in infer_config: - height = infer_config["height"] - else: - height = infer_config.get("default_height", 480) - if "width" in infer_config: - width = infer_config["width"] - else: - width = infer_config.get("default_width", 832) - return { - "height": height, - "width": width, - "num_frames": infer_config.get("num_frames", 81), - "num_inference_steps": infer_config.get("num_inference_steps", 50), - "guidance_scale": infer_config.get("cfg_guidance_scale", 5.0), - } - def _batch_to_list(self, tensor): if tensor.ndim == 4: tensor = tensor.unsqueeze(0) diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index 4ee582a78..d1224a1db 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -74,7 +74,9 @@ def merge(self, other_register): _MODEL_MODULES = { "flux2_dev": "lightx2v_train.model_zoo.flux2.flux2_dev", + "flux2_dev_edit": "lightx2v_train.model_zoo.flux2.flux2_dev_edit", "flux2_klein": "lightx2v_train.model_zoo.flux2.flux2_klein", + "flux2_klein_edit": "lightx2v_train.model_zoo.flux2.flux2_klein_edit", "lingbot_video": "lightx2v_train.model_zoo.wan.lingbot_video", "longcat_image": "lightx2v_train.model_zoo.longcat_image.longcat_image", "longcat_image_edit": "lightx2v_train.model_zoo.longcat_image.longcat_image_edit", @@ -106,7 +108,6 @@ def merge(self, other_register): _INFERENCER_MODULES = { "image_infer": "lightx2v_train.infer.image", - "image_native_infer": "lightx2v_train.infer.image_native", "lingbot_video_t2v_infer": "lightx2v_train.infer.video", "wan_t2v_infer": "lightx2v_train.infer.video", "wan_t2v_14b_infer": "lightx2v_train.infer.video", @@ -119,7 +120,9 @@ def merge(self, other_register): _SAMPLE_PROCESSOR_MODULES = { "flux2_dev": "lightx2v_train.model_zoo.flux2.data_process", + "flux2_dev_edit": "lightx2v_train.model_zoo.flux2.data_process", "flux2_klein": "lightx2v_train.model_zoo.flux2.data_process", + "flux2_klein_edit": "lightx2v_train.model_zoo.flux2.data_process", "longcat_image": "lightx2v_train.model_zoo.longcat_image.data_process", "longcat_image_edit": "lightx2v_train.model_zoo.longcat_image.data_process", "qwen_image": "lightx2v_train.model_zoo.qwen_image.data_process", diff --git a/lightx2v_train/tests/test_longcat_image.py b/lightx2v_train/tests/test_longcat_image.py deleted file mode 100644 index 4ff8a60ab..000000000 --- a/lightx2v_train/tests/test_longcat_image.py +++ /dev/null @@ -1,78 +0,0 @@ -from types import SimpleNamespace - -import torch -from PIL import Image - -from lightx2v_train.model_zoo.longcat_image.data_process import LongCatImageEditDataProcessor -from lightx2v_train.model_zoo.longcat_image.longcat_image import LongCatImageModel -from lightx2v_train.model_zoo.longcat_image.longcat_image_edit import LongCatImageEditModel - - -def _config(name="longcat_image_edit"): - return { - "model": { - "name": name, - "pretrained_model_name_or_path": "/unused", - "running_dtype": "fp32", - }, - "data": {"processor": {"size_multiple": 16, "target_area": 1024 * 1024}}, - "inference": {"enable_cfg_renorm": True, "cfg_renorm_min": 0.0}, - } - - -def test_edit_processor_builds_aligned_source_conditions(): - processor = LongCatImageEditDataProcessor(_config()) - sample = { - "inputs": { - "target_image": Image.new("RGB", (200, 100)), - "source_images": [Image.new("RGB", (200, 100))], - }, - "conditioning": {"prompt": "edit"}, - "meta": {}, - } - - sample = processor(sample) - inputs = sample["inputs"] - height, width = sample["meta"]["target_height"], sample["meta"]["target_width"] - - assert height % 16 == width % 16 == 0 - assert inputs["target_pixel_values"].shape == (3, height, width) - assert inputs["source_vae_pixel_values"].shape == (3, height, width) - assert inputs["source_condition_image"].shape == (3, height // 2, width // 2) - assert inputs["source_condition_image"].dtype == torch.uint8 - assert processor.infer_target_size(sample, 1, 1) == (height, width) - - -def test_edit_denoiser_input_appends_source_modality(): - model = LongCatImageEditModel(_config()) - model.vae = SimpleNamespace(config=SimpleNamespace(block_out_channels=[1, 1, 1, 1])) - noisy_latent = torch.randn(1, 16, 8, 12) - condition = { - "prompt_embed": torch.zeros(1, 520, 8), - "text_ids": torch.zeros(520, 3), - "source_tokens": torch.randn(1, 24, 64), - "source_height": 8, - "source_width": 12, - } - - denoiser_input = model.prepare_denoiser_input(noisy_latent, condition) - - assert denoiser_input.hidden_states.shape == (1, 48, 64) - assert denoiser_input.target_token_length == 24 - assert torch.all(denoiser_input.img_ids[:24, 0] == 1) - assert torch.all(denoiser_input.img_ids[24:, 0] == 2) - assert torch.all(denoiser_input.img_ids[:, 1:] >= 520) - - -def test_cfg_renorm_is_only_used_for_text_to_image(): - positive = torch.tensor([[[3.0, 4.0]]]) - negative = torch.zeros_like(positive) - - text_to_image = LongCatImageModel(_config("longcat_image")) - image_edit = LongCatImageEditModel(_config()) - - renormed = text_to_image.apply_cfg(positive, negative, 2.0) - plain = image_edit.apply_cfg(positive, negative, 2.0) - - assert torch.allclose(torch.linalg.vector_norm(renormed, dim=-1), torch.tensor([[5.0]])) - assert torch.equal(plain, positive * 2) From fdc171718a2afdcad622d0942a2e8c94b305033a Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Sun, 23 Aug 2026 18:30:44 +0000 Subject: [PATCH 21/34] update comment --- lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py | 4 ++++ .../lightx2v_train/model_zoo/flux2/flux2_dev_edit.py | 4 ++++ lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py | 5 ++++- .../lightx2v_train/model_zoo/flux2/flux2_klein_edit.py | 5 +++++ .../lightx2v_train/model_zoo/longcat_image/longcat_image.py | 5 +++++ .../model_zoo/longcat_image/longcat_image_edit.py | 4 ++++ 6 files changed, 26 insertions(+), 1 deletion(-) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index e3ffd92d9..e205da330 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -12,6 +12,10 @@ @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) 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 index f68bd5def..8762562dd 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py @@ -6,6 +6,10 @@ @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 diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index 0b16a29ea..ed7289a17 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -17,7 +17,10 @@ @MODEL_REGISTER("flux2_klein") class Flux2KleinModel(Flux2ModelBase): - """Flux2 Klein Base model. Step-distilled Klein checkpoints are unsupported.""" + """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 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 index aef8d2dd0..2ca39cdcf 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py @@ -6,6 +6,11 @@ @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 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 index 1c8bcb7ae..9075f61a5 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -26,6 +26,11 @@ class LongCatImageDenoiserInput: @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 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 index d90210cb1..149ccfd60 100644 --- 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 @@ -10,6 +10,10 @@ @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 From 474b06d5f3001bd74cdf58a35713807d83151894 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 16:34:38 +0000 Subject: [PATCH 22/34] Add training cache --- lightx2v_train/cache_data.py | 64 +++++ lightx2v_train/data_process/README.md | 59 ++++ lightx2v_train/infer.py | 12 +- .../lightx2v_train/data/image_dataset.py | 93 ++++++- .../lightx2v_train/data/training_cache.py | 48 ++++ .../model_capabilities/__init__.py | 2 + .../model_capabilities/training.py | 24 +- .../lightx2v_train/model_zoo/__init__.py | 10 +- .../lightx2v_train/model_zoo/base.py | 52 +++- .../model_zoo/capability_adapters/common.py | 252 +++++++++++++++++- .../flux2_dopsd_capability.py | 74 +++++ ...2_edit_distribution_matching_capability.py | 7 + .../lightx2v_train/model_zoo/flux2/common.py | 63 +++-- .../model_zoo/flux2/data_process.py | 1 + .../model_zoo/flux2/flux2_dev.py | 4 +- .../model_zoo/flux2/flux2_dev_edit.py | 1 + .../model_zoo/flux2/flux2_klein_edit.py | 1 + ...e_edit_distribution_matching_capability.py | 7 + .../model_zoo/longcat_image/data_process.py | 5 + .../model_zoo/longcat_image/longcat_image.py | 57 ++-- .../longcat_image/longcat_image_edit.py | 5 + .../lightx2v_train/model_zoo/ltx/ltx_t2av.py | 45 ++-- .../model_zoo/minimax_h3/minimax_h3_t2av.py | 26 +- ...e_edit_distribution_matching_capability.py | 7 + .../model_zoo/qwen_image/data_process.py | 1 + .../model_zoo/qwen_image/qwen_image.py | 66 +++-- .../model_zoo/qwen_image/qwen_image_edit.py | 4 + .../model_zoo/wan/lingbot_video.py | 36 +-- .../model_zoo/wan/wan_fastwam.py | 12 +- .../lightx2v_train/model_zoo/wan/wan_t2v.py | 45 ++-- .../model_zoo/wan/wan_ti2v_5b.py | 36 +-- .../lightx2v_train/trainers/__init__.py | 4 + .../trainers/consistency/trainer.py | 21 +- .../lightx2v_train/trainers/dmd/runtime.py | 12 +- .../lightx2v_train/trainers/dmd/trainer.py | 6 +- .../trainers/phased_dmd/roles.py | 30 ++- .../lightx2v_train/trainers/training_cache.py | 156 +++++++++++ .../lightx2v_train/utils/registry.py | 21 +- lightx2v_train/train.py | 12 +- 39 files changed, 1184 insertions(+), 197 deletions(-) create mode 100644 lightx2v_train/cache_data.py create mode 100644 lightx2v_train/lightx2v_train/data/training_cache.py create mode 100644 lightx2v_train/lightx2v_train/trainers/training_cache.py diff --git a/lightx2v_train/cache_data.py b/lightx2v_train/cache_data.py new file mode 100644 index 000000000..a0720f21a --- /dev/null +++ b/lightx2v_train/cache_data.py @@ -0,0 +1,64 @@ +import argparse + +import torch +from lightx2v_train.data import build_data, build_sample_processor, prepare_data +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 +from loguru import logger + + +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("--split", default="train") + parser.add_argument("--save-dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-samples", type=int) + 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, + "max_samples": args.max_samples, + "overwrite": args.overwrite, + } + config["data"][args.split]["preserve_records"] = True + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + init_distributed(config) + setup_logger(config) + + try: + prepare_data(config) + sample_processor = build_sample_processor(config) + dataloader = build_data(config, train_or_val=args.split, 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/data_process/README.md b/lightx2v_train/data_process/README.md index 7394dae6c..f70d93246 100644 --- a/lightx2v_train/data_process/README.md +++ b/lightx2v_train/data_process/README.md @@ -8,6 +8,65 @@ data_process/ └── ltx/build_latent_dataset.py ``` +## 图像训练缓存 + +Qwen-Image、LongCat-Image 和 Flux2 的生成/编辑训练可以预先缓存 VAE latent +与文本、参考图条件。构建缓存时只加载 VAE 和条件编码器，不加载 DiT： + +```bash +cd /path/to/LightX2V + +python lightx2v_train/cache_data.py \ + --config lightx2v_train/configs/train/flow/qwen_image_lora.yaml \ + --output-dir /path/to/qwen_image_cache \ + --save-dtype bf16 +``` + +输出格式对所有图像模型保持一致： + +```text +qwen_image_cache/ +├── cache_data.jsonl +└── cache/ + ├── 00000000.pt + └── 00000001.pt +``` + +`cache_data.jsonl` 保留原始记录，仅增加 `training_cache`： + +```json +{"prompt":"A cat.","target_image":"images/cat.png","training_cache":"cache/00000000.pt"} +``` + +训练时把训练集改为 `image_dataset` 并指向该文件： + +```yaml +data: + use_training_cache: true + + train: + name: image_dataset + data_path: [/path/to/qwen_image_cache/cache_data.jsonl] + num_workers: 8 + persistent_workers: true + prefetch_factor: 2 + pin_memory: true + +inference: + method: none + infer_every_iters: null +``` + +Dataset 会自动读取 PT，并跳过图片解码与模型侧数据处理；训练启动时也只加载 +DiT。Flow Matching、Consistency、DMD 和 DOPSD 通过各自的训练能力接口生成并 +消费统一格式的缓存。DOPSD 缓存训练还需将 +`training.dopsd.trajectory_every_iters` 设为 `null`，因为轨迹可视化需要 VAE 解码。 +缓存按算法需要保存目标 latent 以及正向、无条件或负向条件；修改 prompt、负向 +prompt、图片处理参数或预训练模型后应重新构建。物理 batch 仍固定为 1；多 +worker、持久 worker 与预取用于隐藏单样本 PT 的读取开销。 + +缓存会校验模型、算法和预处理签名；旧版 schema 1 缓存需要重新构建。 + ## 数据集选择 | 训练方式 | 数据集 | 输入 | diff --git a/lightx2v_train/infer.py b/lightx2v_train/infer.py index 25d561389..ef8308933 100644 --- a/lightx2v_train/infer.py +++ b/lightx2v_train/infer.py @@ -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"], @@ -95,7 +97,11 @@ def main(): 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", {}), diff --git a/lightx2v_train/lightx2v_train/data/image_dataset.py b/lightx2v_train/lightx2v_train/data/image_dataset.py index cf74a2047..7b7ae1561 100644 --- a/lightx2v_train/lightx2v_train/data/image_dataset.py +++ b/lightx2v_train/lightx2v_train/data/image_dataset.py @@ -2,10 +2,12 @@ import random from pathlib import Path +import torch from PIL import Image 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 @@ -18,25 +20,42 @@ def __init__( metadata_paths, prompt_dropout_rate=0.0, sample_processor=None, + preserve_records=False, + use_training_cache=False, + unconditional_prompt=" ", + expected_cache_info=None, ): self.prompt_dropout_rate = prompt_dropout_rate self.sample_processor = sample_processor + self.preserve_records = preserve_records + self.uses_training_cache = use_training_cache + self.unconditional_prompt = unconditional_prompt + self.expected_cache_info = expected_cache_info 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.uses_training_cache and any(record["training_cache"] is None for record in self.samples): + raise ValueError("data.use_training_cache=true requires every image record to include training_cache.") + if not self.uses_training_cache and self.sample_processor is None: + raise ValueError("Raw image records require a sample_processor.") def __len__(self): return len(self.samples) def __getitem__(self, index): record = self.samples[index] + if self.uses_training_cache: + return self._load_training_cache(record) prompt = record["prompt"] if random.random() < self.prompt_dropout_rate: - prompt = " " - return self.load_sample(record, prompt=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 = {} @@ -61,6 +80,38 @@ def load_sample(self, record, prompt=None): } return sample if self.sample_processor is None else self.sample_processor(sample) + def _load_training_cache(self, record): + path = record["training_cache"] + cache = self._read_training_cache(path, self.expected_cache_info) + conditioning = cache["conditioning"] + if "positive" not in conditioning: + raise ValueError(f"Invalid training cache at {path}: conditioning.positive is missing.") + if conditioning.get("prompt") != record["prompt"]: + raise ValueError(f"Prompt in {path} does not match cache_data.jsonl. Rebuild the training cache.") + use_unconditional = random.random() < self.prompt_dropout_rate + if use_unconditional and "unconditional" not in conditioning: + raise ValueError(f"Training 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(path) + return cache + + @staticmethod + def _read_training_cache(path, expected_cache_info=None): + cache = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(cache, dict) or not all(key in cache for key in ("inputs", "conditioning", "meta")): + raise ValueError(f"Invalid training cache at {path}: expected inputs, conditioning, and meta 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 expected_cache_info is not None: + for key, expected in 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}.") + return cache + def _load_metadata_samples(self, metadata_path, data_dir): if metadata_path.suffix != ".jsonl": raise ValueError(f"Only metadata list files (.jsonl) are supported, not {metadata_path.suffix}: {metadata_path}") @@ -79,14 +130,21 @@ def _normalize_record(self, record, data_dir): raise ValueError("Each metadata record must include prompt.") source_images = record.get("source_images", []) + training_cache = record.get("training_cache") + if training_cache is not None and not str(training_cache).strip(): + training_cache = None - 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"), + "training_cache": self._resolve_path(training_cache, data_dir) if training_cache is not None else None, } + if self.preserve_records: + normalized["_original_record"] = dict(record) + return normalized def _resolve_path(self, path, data_dir): path = Path(path) @@ -101,10 +159,14 @@ def load_image(image_path): @DATA_REGISTER("image_dataset") -def build_image_dataset(data_config_split, train_or_val="train", sample_processor=None): - if sample_processor is None: - raise ValueError("image_dataset requires a sample_processor") - +def build_image_dataset( + data_config_split, + train_or_val="train", + sample_processor=None, + use_training_cache=False, + unconditional_prompt=" ", + expected_cache_info=None, +): 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" @@ -116,13 +178,30 @@ def build_image_dataset(data_config_split, train_or_val="train", sample_processo metadata_paths=[Path(p) for p in data_path], prompt_dropout_rate=prompt_dropout_rate, sample_processor=sample_processor, + preserve_records=data_config_split.get("preserve_records", False), + use_training_cache=use_training_cache, + unconditional_prompt=unconditional_prompt, + expected_cache_info=expected_cache_info, ) 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", dataset.uses_training_cache) + 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", dataset.uses_training_cache), + collate_fn=_single_sample_collate if dataset.uses_training_cache else None, + **loader_kwargs, ) + + +def _single_sample_collate(samples): + if len(samples) != 1: + raise ValueError(f"Cached image training requires batch_size=1, got {len(samples)} samples.") + return samples[0] 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..9fe0c524d --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/training_cache.py @@ -0,0 +1,48 @@ +import hashlib +import json + +CACHE_SCHEMA_VERSION = 2 + +_RUNTIME_DATA_KEYS = { + "data_path", + "name", + "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"] + signature_data = { + "model": model, + "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(method, {}), + "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/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py index ff3da19d8..50fd08241 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -23,6 +23,7 @@ TeacherForcingCapability, TeacherForcingStepContext, TrainableModelCapability, + TrainingCachePayload, WorldActionTrainingCapability, ) @@ -48,6 +49,7 @@ "SFTStepContext", "TeacherForcingCapability", "TeacherForcingStepContext", + "TrainingCachePayload", "TrainableModelCapability", "WorldActionTrainingCapability", ] diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index 25d0b04be..ea5fe9da2 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -4,7 +4,7 @@ from collections.abc import Collection from contextlib import AbstractContextManager from dataclasses import dataclass, field -from typing import Any, Callable, Iterable, Mapping +from typing import Any, Callable, Iterable, Mapping, TypedDict import torch from torch import Tensor @@ -18,6 +18,12 @@ class LossResult: 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 @@ -26,6 +32,10 @@ class SFTStepContext: 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, @@ -55,6 +65,10 @@ def compute_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.""" @@ -115,6 +129,10 @@ class DistributionMatchingCapability(ModelCapability): 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: @@ -280,6 +298,10 @@ class DopsdCapability(ModelCapability): capability so a new model only needs to describe its own data path. """ + def encode_training_cache(self, batch: Mapping[str, Any]) -> TrainingCachePayload: + """Encode the static sample data consumed by DOPSD training.""" + raise NotImplementedError(f"{type(self).__name__} does not support training-cache encoding.") + @property @abstractmethod def device(self) -> torch.device: diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 8e1b7eace..76804bb5a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -23,14 +23,16 @@ def build_loaded_model( config, *, - transformer_only=False, - reference_model=None, + load_transformer, + load_vae, + load_condition_encoder, ): """Build a model wrapper, load its components, then publish capabilities.""" model = build_model(config) model.load_components( - transformer_only=transformer_only, - reference_model=reference_model, + load_transformer=load_transformer, + load_vae=load_vae, + load_condition_encoder=load_condition_encoder, ) model.ensure_capabilities() return model diff --git a/lightx2v_train/lightx2v_train/model_zoo/base.py b/lightx2v_train/lightx2v_train/model_zoo/base.py index 834bb2703..9c3d08e68 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/base.py +++ b/lightx2v_train/lightx2v_train/model_zoo/base.py @@ -29,12 +29,18 @@ 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( @@ -50,9 +56,23 @@ def register_capabilities(self): 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 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().") @@ -119,6 +139,36 @@ 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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py index efa2b4211..04305bb8c 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/capability_adapters/common.py @@ -1,12 +1,13 @@ from __future__ import annotations import os -from collections.abc import Collection -from typing import Any, Mapping +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, @@ -62,6 +63,157 @@ def _require_singleton_tensor(value, name): 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": @@ -161,6 +313,19 @@ def save_full_model(self, output_path) -> None: 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], @@ -170,8 +335,11 @@ def compute_loss( 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(self.model.encode_to_latent(batch)), + broadcast(latent), "Flow-matching latent", ) noise = broadcast(torch.randn_like(latent, dtype=context.running_dtype)) @@ -186,7 +354,10 @@ def compute_loss( noise, timestep_or_sigma, ) - condition = broadcast(self.model.encode_condition(batch)) + 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, @@ -199,6 +370,26 @@ def compute_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: @@ -212,13 +403,17 @@ 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( - self.model.encode_to_latent(batch), + latent, "Consistency latent", ) def encode_condition(self, batch): - return self.model.encode_condition(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 @@ -255,6 +450,8 @@ def denoiser(self): class GenericDistributionMatchingCapability(BoundCapability, DistributionMatchingCapability): """Distribution-matching operations shared by image flow models.""" + cache_uses_sample_context = False + def __init__( self, model, @@ -272,12 +469,41 @@ def device(self): @property def default_negative_prompt(self): - return None + 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, @@ -338,6 +564,14 @@ def encode_conditions( 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: @@ -487,11 +721,11 @@ def to_dtype(value, dtype): def extract_real_latents(self, batch, dtype, broadcast): with torch.no_grad(): - latent = batch["inputs"].get("latents") + 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: + if latent.ndim == 4 and latent.shape[0] != 1: latent = latent.unsqueeze(0) _require_singleton_tensor(latent, "DMD real latent") return broadcast(latent) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py index 3541c53e4..fdeae2be6 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/capability_adapters/flux2_dopsd_capability.py @@ -12,6 +12,7 @@ DopsdPreparedBatch, DopsdPreparedTeacherBatch, ) +from lightx2v_train.model_zoo.capability_adapters.common import _cached_condition, _training_cache_data, _uses_prompt_dropout @dataclass(frozen=True) @@ -27,6 +28,35 @@ class Flux2DopsdCapability(BoundCapability, DopsdCapability): def device(self) -> torch.device: return self.model.device + def encode_training_cache(self, batch): + from lightx2v_train.trainers.dopsd_core import DopsdConfig + + image = batch.get("inputs", {}).get("target_pixel_values") + if image is None: + raise ValueError("DOPSD cache construction requires inputs.target_pixel_values.") + image = self._image(image, self.model.running_dtype) + latent = self.model.vae.encode(image).latent_dist.mode() + prompt = batch["conditioning"]["prompt"] + config = DopsdConfig.from_training_config(self.model.config["training"]) + teacher_prompts = config.teacher_prompts(prompt) + prompts = { + "positive": prompt, + "teacher": teacher_prompts, + } + contextual_roles = {"positive"} + if _uses_prompt_dropout(self.model): + prompts["unconditional"] = self.model.unconditional_prompt + prompts["teacher_unconditional"] = config.teacher_prompts(self.model.unconditional_prompt) + contextual_roles.add("unconditional") + return _training_cache_data( + self.model, + batch, + inputs={"dopsd_reference_latents": self.model._normalize_patch_latents(latent)}, + prompts=prompts, + contextual_roles=contextual_roles, + conditioning_meta={"teacher_prompt": teacher_prompts}, + ) + def _image(self, image: torch.Tensor, running_dtype: torch.dtype) -> torch.Tensor: if image.ndim == 3: image = image.unsqueeze(0) @@ -56,12 +86,56 @@ def _reference(self, image: torch.Tensor) -> Flux2DopsdReference: ids=reference_ids, ) + def _cached_reference(self, latent: torch.Tensor) -> Flux2DopsdReference: + reference_ids = Flux2KleinPipeline._prepare_image_ids([latent[:1]]) + reference_ids = reference_ids.repeat(latent.shape[0], 1, 1).to(self.device) + return Flux2DopsdReference( + latents=Flux2KleinPipeline._pack_latents(latent), + ids=reference_ids, + ) + + @staticmethod + def _dimension(meta, key): + value = meta[key] + if torch.is_tensor(value): + value = value.reshape(-1)[0].item() + elif isinstance(value, (list, tuple)): + value = value[0] + return int(value) + def prepare_training_batch( self, batch: Mapping[str, Any], teacher_prompts: list[str], running_dtype: torch.dtype, ) -> DopsdPreparedBatch: + cached_latent = batch.get("inputs", {}).get("dopsd_reference_latents") + if cached_latent is not None: + cached_teacher_prompt = batch["conditioning"].get("teacher_prompt") + if list(cached_teacher_prompt or []) != list(teacher_prompts): + raise ValueError("Cached DOPSD teacher prompt does not match the current training configuration. Rebuild the cache.") + cached_latent = cached_latent.to(device=self.device, dtype=running_dtype) + meta = batch.get("meta", {}) + height = self._dimension(meta, "target_height") + width = self._dimension(meta, "target_width") + initial_state, state_ids, latent_hw = self._initial_state(height, width, generator=None) + return DopsdPreparedBatch( + initial_state=initial_state, + state_ids=state_ids, + student_condition=_cached_condition(batch, self.model), + teacher_condition=_cached_condition( + batch, + self.model, + role="teacher_unconditional" if batch["conditioning"].get("active") == "unconditional" else "teacher", + ), + teacher_reference=self._cached_reference(cached_latent), + latent_hw=latent_hw, + ) + + 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.dopsd_reference_latents entry.") + image = batch.get("inputs", {}).get("target_pixel_values") if image is None: raise ValueError("DOPSD requires inputs.target_pixel_values.") 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 index aa0702d86..31a357604 100644 --- 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 @@ -2,6 +2,7 @@ from lightx2v_train.model_zoo.capability_adapters.common import ( GenericDistributionMatchingCapability, + _cached_condition_pair, _negative_prompt, _require_single_prompt, ) @@ -10,10 +11,16 @@ 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)) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py index b005fbfcd..323fb124e 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/common.py @@ -23,46 +23,64 @@ class Flux2ModelBase(BaseModel): requires_source_images = False target_latent_mode = "sample" default_text_encoder_out_layers = () - - def load_components(self, transformer_only=False, reference_model=None): + 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 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._copy_model_state(reference_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) - return + 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.transformer = self.load_transformer(model_path) - - self.text_pipeline.text_encoder.requires_grad_(False) + 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 _copy_model_state(self, reference_model): - del reference_model + 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"] @@ -94,7 +112,9 @@ def fsdp2_shard_plan(self, fsdp_config): @property def vae_scale_factor(self): - return 2 ** (len(self.vae.config.block_out_channels) - 1) + 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) @@ -111,6 +131,12 @@ def _latent_statistics(self, latents): 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) @@ -118,7 +144,7 @@ def encode_to_latent(self, sample): 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, self.target_latent_mode)() + latent = getattr(distribution, mode)() return self._normalize_patch_latents(latent) def encode_condition(self, sample): @@ -127,6 +153,11 @@ def encode_condition(self, sample): 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( diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py index 2f9fc5653..c88292c54 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -56,6 +56,7 @@ 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"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py index e205da330..81414648b 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev.py @@ -39,8 +39,8 @@ def register_capabilities(self): def _configure_model(self): self.guidance_scale = float(self.config["model"].get("guidance_scale", 4.0)) - def _copy_model_state(self, reference_model): - self.guidance_scale = reference_model.guidance_scale + 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( 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 index 8762562dd..04598a17b 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_dev_edit.py @@ -13,3 +13,4 @@ class Flux2DevEditModel(Flux2DevModel): 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_edit.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py index 2ca39cdcf..b38aff5b3 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein_edit.py @@ -14,4 +14,5 @@ class Flux2KleinEditModel(Flux2KleinModel): requires_source_images = True target_latent_mode = "mode" distribution_matching_capability_cls = Flux2EditDistributionMatchingCapability + shared_condition_keys = ("reference_tokens", "reference_ids") supports_dopsd = False 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 index d6088ad45..2147e57c7 100644 --- 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 @@ -2,16 +2,23 @@ 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)) 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 index 2133da287..5ac13b1da 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/data_process.py @@ -88,6 +88,7 @@ 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"] @@ -117,6 +118,10 @@ def infer_target_size(self, sample, default_height, default_width): @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) 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 index 9075f61a5..4c7ed5907 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/longcat_image/longcat_image.py @@ -47,7 +47,7 @@ def register_capabilities(self): self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( - channels_path="vae.config.latent_channels", + channels_path="latent_channels", ), guidance_in_denoiser_space=True, ), @@ -57,33 +57,51 @@ def register_capabilities(self): LongCatImageConsistencyModelCapability(self), ) - 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 + 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() - return - model_path = self.config["model"]["pretrained_model_name_or_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 = AutoencoderKL.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_config = self.vae.config self.vae.requires_grad_(False) + self.vae.eval() self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) - self._maybe_set_attention_backend() + + 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"] @@ -120,7 +138,9 @@ def fsdp2_shard_plan(self, fsdp_config): @property def vae_scale_factor(self): - return 2 ** (len(self.vae.config.block_out_channels) - 1) + 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) @@ -133,13 +153,20 @@ def _denormalize_latents(self, latents): 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) - return self._normalize_latents(self.vae.encode(image).latent_dist.sample()) + latent = getattr(self.vae.encode(image).latent_dist, mode)() + return self._normalize_latents(latent) def encode_condition(self, sample): prompt = sample["conditioning"]["prompt"] @@ -234,7 +261,7 @@ def apply_cfg(self, positive, negative, guidance_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.vae.config.latent_channels, latent_height, latent_width) + 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): 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 index 149ccfd60..e2eabdb78 100644 --- 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 @@ -18,11 +18,16 @@ class LongCatImageEditModel(LongCatImageModel): 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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py index f9b0a7a7a..343c430d1 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py +++ b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py @@ -71,7 +71,14 @@ def register_capabilities(self): LTXTeacherForcingCapability(self), ) - def load_components(self, transformer_only=False, reference_model=None): + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + del load_vae 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))) @@ -80,9 +87,9 @@ def load_components(self, transformer_only=False, reference_model=None): 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)) + should_load_transformer = load_transformer and bool(model_config.get("load_transformer", True)) + should_load_embeddings_processor = load_condition_encoder and bool(model_config.get("load_embeddings_processor", True)) + should_load_text_encoder = load_condition_encoder and 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") @@ -97,17 +104,8 @@ def load_components(self, transformer_only=False, reference_model=None): 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: + transformer_path = self._resolve_transformer_path(model_path, model_config) + if should_load_transformer: self.transformer = self._load_transformer(transformer_path) self._configure_causal_transformer() else: @@ -115,17 +113,24 @@ def load_components(self, transformer_only=False, reference_model=None): self.embeddings_processor = getattr(self, "embeddings_processor", None) self.text_encoder = getattr(self, "text_encoder", None) - if self.load_embeddings_processor: + if should_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: + if hasattr(self.embeddings_processor, "feature_extractor") and not should_load_text_encoder: self.embeddings_processor.feature_extractor = None - if self.load_text_encoder: + if should_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 reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self.embeddings_processor = source.embeddings_processor + self.text_encoder = getattr(source, "text_encoder", None) + self.text_encoder_path = getattr(source, "text_encoder_path", self.text_encoder_path) + self.text_encoder_cpu = getattr(source, "text_encoder_cpu", self.text_encoder_cpu) + def _load_transformer(self, model_path): builder = SingleGPUModelBuilder( model_path=str(model_path), @@ -136,9 +141,9 @@ def _load_transformer(self, model_path): 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): + def _resolve_transformer_path(self, model_path, model_config): self.loaded_student_checkpoint_path = None - if transformer_only or model_config.get("name") != "ltx_t2av_ar": + if model_config.get("name") != "ltx_t2av_ar": return model_path student_config = model_config.get("student", {}) 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 index 9f79bf5d6..ad1737c33 100644 --- 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 @@ -40,8 +40,14 @@ def register_capabilities(self): MiniMaxH3ConsistencyModelCapability(self), ) - def load_components(self, transformer_only=False, reference_model=None): - del transformer_only, reference_model + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + del 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")) @@ -54,13 +60,15 @@ def load_components(self, transformer_only=False, reference_model=None): if self.video_latent_channels <= 0 or self.audio_latent_channels <= 0 or self.vae_spatial_scale_factor <= 0: raise ValueError("MiniMax-H3 latent channels and VAE spatial scale must be positive.") 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) + self.transformer = None + 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 denoiser_module(self): return self.transformer diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py index a6106bed7..09d1cad9f 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/capability_adapters/qwen_image_edit_distribution_matching_capability.py @@ -4,6 +4,7 @@ from lightx2v_train.model_zoo.capability_adapters.common import ( GenericDistributionMatchingCapability, + _cached_condition_pair, _negative_prompt, _require_single_prompt, ) @@ -12,10 +13,16 @@ class QwenImageEditDistributionMatchingCapability(GenericDistributionMatchingCapability): """Keep source-image conditioning in Qwen-Image-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)) diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py index f8a5c6e74..5966bb40d 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py @@ -70,6 +70,7 @@ class QwenImageDataProcessor: 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"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 50bc4bd6d..22b43a4cf 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -43,7 +43,7 @@ def register_capabilities(self): self.distribution_matching_capability_cls( self, latent_geometry=SpatialLatentGeometry( - channels_path="vae.config.z_dim", + channels_path="latent_channels", temporal_size=1, ), ), @@ -53,30 +53,53 @@ def register_capabilities(self): QwenImageConsistencyModelCapability(self), ) - 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.vae_scale_factor = reference_model.vae_scale_factor - self.image_processor = reference_model.image_processor - self.transformer = self.load_transformer() - return + 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() + + 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.vae = AutoencoderKLQwenImage.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.text_pipeline.text_encoder.eval() + + def _load_vae(self, model_path): + self.vae = AutoencoderKLQwenImage.from_pretrained(model_path, subfolder="vae").to(self.device, dtype=self.running_dtype) + self.vae_config = self.vae.config self.vae.requires_grad_(False) + self.vae.eval() self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + def _load_vae_config(self, model_path): + self.vae_config = AutoencoderKLQwenImage.load_config(model_path, subfolder="vae") + self.vae_scale_factor = 2 ** len(self.vae_config["temperal_downsample"]) + + def reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self.vae_scale_factor = source.vae_scale_factor + + @property + def latent_channels(self): + return int(self.vae_config["z_dim"] if isinstance(self.vae_config, dict) else self.vae_config.z_dim) + def load_transformer(self): model_path = self.config["model"]["pretrained_model_name_or_path"] return QwenImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer").to(self.device, dtype=self.running_dtype) @@ -106,9 +129,10 @@ def fsdp2_shard_plan(self, fsdp_config): ] def _latent_statistics(self, latents): - shape = (1, self.vae.config.z_dim, 1, 1, 1) - mean = latents.new_tensor(self.vae.config.latents_mean).view(shape) - std = latents.new_tensor(self.vae.config.latents_std).view(shape) + config = self.vae.config + shape = (1, config.z_dim, 1, 1, 1) + mean = latents.new_tensor(config.latents_mean).view(shape) + std = latents.new_tensor(config.latents_std).view(shape) return mean, std def _normalize_latents(self, latents): @@ -120,6 +144,12 @@ def _denormalize_latents(self, latents): return latents * std + mean 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) @@ -127,7 +157,7 @@ def encode_to_latent(self, sample): 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) pixel_values = image.unsqueeze(2) - latent = self.vae.encode(pixel_values).latent_dist.sample() # (B, C, T, H, W) + latent = getattr(self.vae.encode(pixel_values).latent_dist, mode)() # (B, C, T, H, W) return self._normalize_latents(latent) def encode_condition(self, sample): @@ -189,7 +219,7 @@ def postprocess_denoiser_output(self, prediction, denoiser_input): def prepare_infer_latents(self, height, width, generator=None): latent_h = height // self.vae_scale_factor latent_w = width // self.vae_scale_factor - shape = (1, self.vae.config.z_dim, 1, latent_h, latent_w) + shape = (1, self.latent_channels, 1, latent_h, latent_w) return torch.randn(shape, generator=generator, device=self.device, dtype=self.running_dtype) def decode_latent(self, latent): diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py index 81cc54d0a..62d21ff0f 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image_edit.py @@ -16,11 +16,15 @@ class QwenImageEditModel(QwenImageModel): pipeline_cls = QwenImageEditPlusPipeline distribution_matching_capability_cls = QwenImageEditDistributionMatchingCapability + shared_condition_keys = ("source_latents", "source_img_shapes") 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_conditions_with_source(self, sample, prompts): inputs = sample["inputs"] condition_tensors = inputs.get("source_condition_images", []) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py index b7b942190..650b47a21 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/lingbot_video.py @@ -64,15 +64,21 @@ def register_capabilities(self): vae_scale_factor_temporal = 4 vae_scale_factor_spatial = 8 - def load_components(self, transformer_only=False, reference_model=None): + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): model_config = self.config["model"] self.model_path = os.path.abspath(os.path.expanduser(str(model_config["pretrained_model_name_or_path"]))) self.transformer_param_dtype = get_running_dtype(model_config.get("transformer_param_dtype", "bf16")) self.text_encoder_dtype = get_running_dtype(model_config.get("text_encoder_dtype", "bf16")) self.vae_dtype = get_running_dtype(model_config.get("vae_dtype", "fp32")) - self.load_transformer = bool(model_config.get("load_transformer", True)) - self.load_text_encoder = bool(model_config.get("load_text_encoder", True)) - self.load_vae = bool(model_config.get("load_vae", False)) + should_load_transformer = load_transformer and bool(model_config.get("load_transformer", True)) + should_load_text_encoder = load_condition_encoder and bool(model_config.get("load_text_encoder", True)) + should_load_vae = load_vae and bool(model_config.get("load_vae", False)) self.text_encoder_cpu = bool(model_config.get("text_encoder_cpu", True)) self.max_sequence_length = int(model_config.get("max_sequence_length", TOKEN_LENGTH)) self.hidden_state_skip_layer = int(model_config.get("hidden_state_skip_layer", HIDDEN_STATE_SKIP_LAYER)) @@ -82,21 +88,19 @@ def load_components(self, transformer_only=False, reference_model=None): self.text_encoder = None self.processor = None self.vae = None - if transformer_only: - self.load_text_encoder = False - self.load_vae = False - if reference_model is not None: - self.text_encoder = reference_model.text_encoder - self.processor = reference_model.processor - self.vae = reference_model.vae - self._crop_start = reference_model._crop_start - - self.transformer = self._load_transformer() if self.load_transformer else None - if self.load_text_encoder: + + self.transformer = self._load_transformer() if should_load_transformer else None + if should_load_text_encoder: self.text_encoder, self.processor = self._load_text_components() - if self.load_vae: + if should_load_vae: self.vae = self._load_vae() + def reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self.text_encoder = source.text_encoder + self.processor = source.processor + self._crop_start = source._crop_start + def _load_transformer(self): try: from ..native.lingbot_video import LingBotVideoTransformer3DModel diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py index 517fe5d81..67a7af58a 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py @@ -43,9 +43,15 @@ def register_capabilities(self): WanWorldActionCapability(self), ) - def load_components(self, transformer_only=False, reference_model=None): - if transformer_only: - raise ValueError("wan_fastwam does not support transformer_only loading.") + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): + if not (load_transformer and load_vae and load_condition_encoder): + raise ValueError("wan_fastwam does not support selective component loading.") model_path = _resolve_local_path(self.model_config.get("model_path"), "model_path", directory=True) configured_action_dit_path = self.model_config.get("action_dit_pretrained_path") diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py index bd16f02a2..e0c914046 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_t2v.py @@ -76,13 +76,19 @@ def register_capabilities(self): WanTeacherForcingCapability(self), ) - def load_components(self, transformer_only=False, reference_model=None): + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): model_config = self.config["model"] model_path = model_config["pretrained_model_name_or_path"] - self.load_vae = model_config.get("load_vae", True) - self.load_text_encoder = model_config.get("load_text_encoder", True) - self.load_transformer = model_config.get("load_transformer", True) + should_load_vae = load_vae and model_config.get("load_vae", True) + should_load_text_encoder = load_condition_encoder and model_config.get("load_text_encoder", True) + should_load_transformer = load_transformer and model_config.get("load_transformer", True) training_config = self.config.get("training", {}) teacher_forcing_config = training_config.get("teacher_forcing", {}) model_name = model_config.get("name") @@ -124,34 +130,18 @@ def load_components(self, transformer_only=False, reference_model=None): self.text_encoder = None self.text_pipeline = None - if transformer_only: - if reference_model is not None: - self.vae = reference_model.vae - self.text_encoder = reference_model.text_encoder - self.text_pipeline = reference_model.text_pipeline - self.vae_stride = reference_model.vae_stride - self.patch_size = reference_model.patch_size - self.max_sequence_length = reference_model.max_sequence_length - self.vae_scale_factor_temporal = reference_model.vae_scale_factor_temporal - self.vae_scale_factor_spatial = reference_model.vae_scale_factor_spatial - self.transformer = self._load_transformer(model_path) - self._configure_transformer() - self.vae_scale_factor_temporal = self.vae_stride[0] - self.vae_scale_factor_spatial = self.vae_stride[1] - return - - if self.load_transformer: + if should_load_transformer: self.transformer = self._load_transformer(model_path) self._configure_transformer() else: self.transformer = None - if self.load_vae: + if should_load_vae: vae_checkpoint = os.path.join(model_path, "Wan2.1_VAE.pth") self.vae = WanVAE(vae_pth=vae_checkpoint, dtype=self.vae_dtype, device=self.device) self.vae.model.requires_grad_(False) - if self.load_text_encoder: + if should_load_text_encoder: t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth") t5_tokenizer = os.path.join(model_path, "google/umt5-xxl") self.text_encoder = T5EncoderModel( @@ -168,6 +158,15 @@ def load_components(self, transformer_only=False, reference_model=None): self.vae_scale_factor_temporal = self.vae_stride[0] self.vae_scale_factor_spatial = self.vae_stride[1] + def reuse_frozen_components_from(self, source): + super().reuse_frozen_components_from(source) + self.text_encoder = source.text_encoder + self.vae_stride = source.vae_stride + self.patch_size = source.patch_size + self.max_sequence_length = source.max_sequence_length + self.vae_scale_factor_temporal = source.vae_scale_factor_temporal + self.vae_scale_factor_spatial = source.vae_scale_factor_spatial + def _load_transformer(self, model_path): if self.use_causal_transformer: transformer = CausalWanModel.from_pretrained( diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py index cd90e97fb..7b62675c9 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_ti2v_5b.py @@ -16,13 +16,19 @@ @MODEL_REGISTER("wan_ti2v_5b_ar") @MODEL_REGISTER("wan_ti2v_5b") class WanTI2V5BModel(WanT2VModel): - def load_components(self, transformer_only=False, reference_model=None): + def load_components( + self, + *, + load_transformer, + load_vae, + load_condition_encoder, + ): model_config = self.config["model"] model_path = model_config["pretrained_model_name_or_path"] - self.load_vae = model_config.get("load_vae", True) - self.load_text_encoder = model_config.get("load_text_encoder", True) - self.load_transformer = model_config.get("load_transformer", True) + should_load_vae = load_vae and model_config.get("load_vae", True) + should_load_text_encoder = load_condition_encoder and model_config.get("load_text_encoder", True) + should_load_transformer = load_transformer and model_config.get("load_transformer", True) training_config = self.config.get("training", {}) teacher_forcing_config = training_config.get("teacher_forcing", {}) model_name = model_config.get("name") @@ -63,34 +69,18 @@ def load_components(self, transformer_only=False, reference_model=None): self.text_encoder = None self.text_pipeline = None - if transformer_only: - if reference_model is not None: - self.vae = reference_model.vae - self.text_encoder = reference_model.text_encoder - self.text_pipeline = reference_model.text_pipeline - self.vae_stride = reference_model.vae_stride - self.patch_size = reference_model.patch_size - self.max_sequence_length = reference_model.max_sequence_length - self.vae_scale_factor_temporal = reference_model.vae_scale_factor_temporal - self.vae_scale_factor_spatial = reference_model.vae_scale_factor_spatial - self.transformer = self._load_transformer(model_path) - self._configure_transformer() - self.vae_scale_factor_temporal = self.vae_stride[0] - self.vae_scale_factor_spatial = self.vae_stride[1] - return - - if self.load_transformer: + if should_load_transformer: self.transformer = self._load_transformer(model_path) self._configure_transformer() else: self.transformer = None - if self.load_vae: + if should_load_vae: vae_checkpoint = model_config.get("vae_checkpoint", os.path.join(model_path, "Wan2.2_VAE.pth")) self.vae = Wan2_2_VAE(vae_pth=vae_checkpoint, dtype=self.vae_dtype, device=self.device) self.vae.model.requires_grad_(False) - if self.load_text_encoder: + if should_load_text_encoder: t5_checkpoint = model_config.get("t5_checkpoint", os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth")) t5_tokenizer = model_config.get("t5_tokenizer", os.path.join(model_path, "google/umt5-xxl")) self.text_encoder = T5EncoderModel( diff --git a/lightx2v_train/lightx2v_train/trainers/__init__.py b/lightx2v_train/lightx2v_train/trainers/__init__.py index 225a8cab4..75aa1d806 100644 --- a/lightx2v_train/lightx2v_train/trainers/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/__init__.py @@ -26,6 +26,10 @@ ".teacher_forcing", "TeacherForcingTrainer", ), + "TrainingCacheTrainer": ( + ".training_cache", + "TrainingCacheTrainer", + ), "WorldActionTrainer": (".world_action", "WorldActionTrainer"), } diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py index 9c0365632..21c3b445f 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/trainer.py @@ -134,9 +134,11 @@ def _build_frozen_teacher(self): teacher_config["model"].update(copy.deepcopy(teacher_override)) teacher_model = build_loaded_model( teacher_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + teacher_model.reuse_frozen_components_from(self.model) self.teacher_consistency_model = teacher_model.capabilities.require(ConsistencyModelCapability) self.teacher_consistency_model.set_frozen() teacher_model.capabilities.require(ParallelCapability).apply(self.config) @@ -172,9 +174,11 @@ def _build_frozen_reference(self, spec): reference_config["model"] = model_config reference_model = build_loaded_model( reference_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + reference_model.reuse_frozen_components_from(self.model) capability = reference_model.capabilities.require(ConsistencyModelCapability) if self.train_type == "lora": reference_model.capabilities.require(TrainableModelCapability).configure( @@ -251,7 +255,9 @@ def _encode_negative_condition(self, sample): if len(values) != 1: raise ValueError(f"Expected exactly one negative prompt, got {len(values)}.") - fallback = self.objective.negative_prompt or " " + fallback = self.objective.negative_prompt + if not isinstance(fallback, str): + fallback = self.model.unconditional_prompt values = [value if isinstance(value, str) and value.strip() else fallback for value in values] encoded_prompt = values[0] if isinstance(prompt, str) else values @@ -259,8 +265,11 @@ def _encode_negative_condition(self, sample): negative_sample["conditioning"] = dict(conditioning) cached_negative = negative_sample["conditioning"].get("negative") if cached_negative is None: + cache_path = sample.get("meta", {}).get("training_cache_path") + if cache_path is not None: + raise KeyError(f"Training cache {cache_path} has no conditioning.negative entry required by this consistency objective.") negative_sample["conditioning"].pop("positive", None) else: - negative_sample["conditioning"]["positive"] = cached_negative + negative_sample["conditioning"]["active"] = "negative" negative_sample["conditioning"]["prompt"] = encoded_prompt return self.consistency_model.encode_condition(negative_sample) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py index 711a33537..e861284bf 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/runtime.py @@ -184,9 +184,11 @@ def setup(self, resume_ckpt_path=None): self.fake_model_config = copy.deepcopy(fake_model_config) self.fake_model = build_loaded_model( fake_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.fake_model.reuse_frozen_components_from(self.model) self.fake = self.fake_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model(self.fake_model, role="fake") self.fake_model.capabilities.require(ParallelCapability).apply(self.config) @@ -201,9 +203,11 @@ def setup(self, resume_ckpt_path=None): teacher_model_config["model"].update(copy.deepcopy(self.model_config["teacher"])) self.teacher_model = build_loaded_model( teacher_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.teacher_model.reuse_frozen_components_from(self.model) self.teacher = self.teacher_model.capabilities.require(DistributionMatchingCapability) self.teacher.denoiser().requires_grad_(False) self.teacher.set_training(False) diff --git a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py index 30af6db9d..9d05f9be6 100644 --- a/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py +++ b/lightx2v_train/lightx2v_train/trainers/dmd/trainer.py @@ -170,9 +170,11 @@ def _setup_fake_real_resources(self): model_config = copy.deepcopy(self.fake_model_config) self.fake_real_model = build_loaded_model( model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.fake_real_model.reuse_frozen_components_from(self.model) self.fake_real = self.fake_real_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( self.fake_real_model, diff --git a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py index 722148c6d..1bda39b09 100644 --- a/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py +++ b/lightx2v_train/lightx2v_train/trainers/phased_dmd/roles.py @@ -245,9 +245,11 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): student_2_model_config = self._role_model_config("student_2") self.student_2_model = build_loaded_model( student_2_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.student_2_model.reuse_frozen_components_from(self.model) self.student_2 = self.student_2_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( self.student_2_model, @@ -260,9 +262,11 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): fake_2_model_config = self._role_model_config("fake_2") self.fake_2_model = build_loaded_model( fake_2_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.fake_2_model.reuse_frozen_components_from(self.model) self.fake_2 = self.fake_2_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model(self.fake_2_model, role="fake_2") self.fake_2_model.capabilities.require(ParallelCapability).apply(self.config) @@ -276,9 +280,11 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): fake_low_high_model_config = self._role_model_config("teacher") self.fake_low_high_model = build_loaded_model( fake_low_high_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.fake_low_high_model.reuse_frozen_components_from(self.model) fake_low_high_model_path = fake_low_high_model_config["model"]["pretrained_model_name_or_path"] self.fake_low_high = self.fake_low_high_model.capabilities.require(DistributionMatchingCapability) self._setup_trainable_model( @@ -300,9 +306,11 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): role_config = self._role_model_config(source_role) model = build_loaded_model( role_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + model.reuse_frozen_components_from(self.model) self._setup_trainable_model(model, role=role) model.capabilities.require(ParallelCapability).apply(self.config) if self.gradient_checkpointing: @@ -333,9 +341,11 @@ def setup_resources(self, resume_ckpt_path=None, base_setup=None): teacher_2_model_config = self._role_model_config("teacher_2") self.teacher_2_model = build_loaded_model( teacher_2_model_config, - transformer_only=True, - reference_model=self.model, + load_transformer=True, + load_vae=False, + load_condition_encoder=False, ) + self.teacher_2_model.reuse_frozen_components_from(self.model) self.teacher_2 = self.teacher_2_model.capabilities.require(DistributionMatchingCapability) self.teacher_2.denoiser().requires_grad_(False) self.teacher_2.set_training(False) diff --git a/lightx2v_train/lightx2v_train/trainers/training_cache.py b/lightx2v_train/lightx2v_train/trainers/training_cache.py new file mode 100644 index 000000000..f80dd4e29 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/training_cache.py @@ -0,0 +1,156 @@ +import json +import os +from pathlib import Path + +import torch +import torch.distributed as dist +from loguru import logger + +from lightx2v_train.data.training_cache import preserve_cache_dtype, training_cache_info +from lightx2v_train.data.utils import require_singleton_dataloader +from lightx2v_train.model_capabilities import ( + ConsistencyModelCapability, + DistributionMatchingCapability, + DopsdCapability, + FlowMatchingSFTCapability, +) +from lightx2v_train.runtime.distributed import get_rank, get_sequence_parallel_rank, is_distributed, is_main_process +from lightx2v_train.utils.registry import TRAINER_REGISTER + +CACHE_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +CACHE_CAPABILITIES = { + "flow_matching": FlowMatchingSFTCapability, + "consistency": ConsistencyModelCapability, + "dmd": DistributionMatchingCapability, + "dopsd": DopsdCapability, +} + + +def _to_cpu(value, dtype, key=None): + if torch.is_tensor(value): + value = value.detach().cpu().contiguous() + return value.to(dtype) if value.is_floating_point() and not preserve_cache_dtype(key) else value + if isinstance(value, dict): + return {name: _to_cpu(item, dtype, name) for name, item in value.items()} + if isinstance(value, tuple): + return tuple(_to_cpu(item, dtype, key) for item in value) + if isinstance(value, list): + return [_to_cpu(item, dtype, key) for item in value] + return value + + +def _atomic_save(value, path): + temporary = path.with_suffix(f"{path.suffix}.rank{get_rank():05d}.tmp") + torch.save(value, temporary) + os.replace(temporary, path) + + +def _write_jsonl(records, path): + temporary = path.with_suffix(f"{path.suffix}.tmp") + with temporary.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + os.replace(temporary, path) + + +def _dataset_index(sample): + value = sample["meta"].pop("dataset_index") + return int(value.item() if torch.is_tensor(value) else value) + + +@TRAINER_REGISTER("training_cache") +class TrainingCacheTrainer: + def __init__(self, config): + self.config = config + self.cache_config = config["training_cache"] + self.training_method = config["training"]["method"] + self.cache_info = { + **training_cache_info(config), + "storage_dtype": self.cache_config["save_dtype"], + } + + def set_model(self, model): + capability_type = CACHE_CAPABILITIES.get(self.training_method) + if capability_type is None: + supported = ", ".join(sorted(CACHE_CAPABILITIES)) + raise ValueError(f"Training cache does not support {self.training_method!r}; expected one of: {supported}.") + self.model = model + self.encoder = model.ensure_capabilities().require(capability_type) + + def set_data(self, dataloader_train, dataloader_eval=None): + del dataloader_eval + require_singleton_dataloader(dataloader_train, "Cache dataloader") + self.dataloader = dataloader_train + + def _encode(self, sample, dtype): + encoded = self.encoder.encode_training_cache(sample) + cache = { + "cache_info": self.cache_info, + **encoded, + } + self._validate(cache, sample["conditioning"]["prompt"]) + return _to_cpu(cache, dtype) + + def _validate(self, cache, prompt, path=""): + 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 cache.get("cache_info") != self.cache_info: + raise ValueError(f"Training cache at {path} is incompatible with the current configuration. Rebuild it.") + conditioning = cache["conditioning"] + if not isinstance(conditioning, dict) or "positive" not in conditioning: + raise ValueError(f"Invalid training cache at {path}: conditioning.positive is missing.") + if conditioning.get("prompt") != prompt: + raise ValueError(f"Training cache at {path} has a different prompt. Rebuild it.") + + def _gather_records(self, records): + if not is_distributed(): + return records + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, records) + return [item for rank_records in gathered for item in rank_records] + + @torch.inference_mode() + def train(self): + output_dir = Path(self.cache_config["output_dir"]).resolve() + cache_dir = output_dir / "cache" + cache_dir.mkdir(parents=True, exist_ok=True) + + dataset = self.dataloader.dataset + max_samples = self.cache_config["max_samples"] + sample_count = len(dataset) if max_samples is None else min(len(dataset), max_samples) + dtype = CACHE_DTYPES[self.cache_config["save_dtype"]] + records = [] + + for sample in self.dataloader: + index = _dataset_index(sample) + if get_sequence_parallel_rank() != 0 or index >= sample_count: + continue + + record = dict(dataset.samples[index]["_original_record"]) + sample["conditioning"]["prompt"] = dataset.samples[index]["prompt"] + cache_path = cache_dir / f"{index:08d}.pt" + if self.cache_config["overwrite"] or not cache_path.exists(): + torch.manual_seed(self.cache_config["seed"] + index) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(self.cache_config["seed"] + index) + _atomic_save(self._encode(sample, dtype), cache_path) + else: + existing = torch.load(cache_path, map_location="cpu", weights_only=True) + self._validate(existing, sample["conditioning"]["prompt"], cache_path) + + record["training_cache"] = cache_path.relative_to(output_dir).as_posix() + records.append((index, record)) + logger.info("[cache] {}/{} -> {}", index + 1, sample_count, cache_path) + + records = self._gather_records(records) + if is_main_process(): + indexed_records = {index: record for index, record in records} + if sorted(indexed_records) != list(range(sample_count)): + raise RuntimeError("Training cache is incomplete.") + _write_jsonl([indexed_records[index] for index in range(sample_count)], output_dir / "cache_data.jsonl") + logger.info("[cache] wrote {} samples to {}", sample_count, output_dir / "cache_data.jsonl") diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index d1224a1db..a5047acec 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -103,6 +103,7 @@ def merge(self, other_register): "phased_dmd": "lightx2v_train.trainers.phased_dmd.trainer", "sgmd": "lightx2v_train.trainers.sgmd", "teacher_forcing": "lightx2v_train.trainers.teacher_forcing", + "training_cache": "lightx2v_train.trainers.training_cache", "world_action": "lightx2v_train.trainers.world_action", } @@ -159,7 +160,7 @@ def build_model(config): def build_trainer(config): - name = config["training"]["method"] + name = "training_cache" if "training_cache" in config else config["training"]["method"] _ensure_registered(name, TRAINER_REGISTER, _TRAINER_MODULES) if name not in TRAINER_REGISTER: available = ", ".join(sorted(TRAINER_REGISTER.keys())) @@ -189,12 +190,19 @@ def build_sample_processor(config): def build_data(config, train_or_val, sample_processor=None): + from lightx2v_train.data.training_cache import training_cache_info + data_config = config.get("data", {}) if train_or_val not in data_config: available_splits = ", ".join(repr(k) for k in sorted(data_config.keys())) raise ValueError(f"config['data'] has no key {train_or_val!r}. Available keys: {available_splits}") data_config_split = dict(data_config[train_or_val]) + if "training_cache" in config: + data_config_split["prompt_dropout_rate"] = 0.0 data_name = data_config_split.get("name", "image_dataset") + use_training_cache = data_config.get("use_training_cache", False) if train_or_val == "train" and "training_cache" not in config else False + if use_training_cache and data_name != "image_dataset": + raise ValueError("data.use_training_cache=true requires data.train.name=image_dataset.") if train_or_val == "train" and data_name == "prompt_dataset": image_sizes = config.get("training", {}).get("dmd", {}).get("image_sizes") if image_sizes is not None: @@ -203,5 +211,14 @@ def build_data(config, train_or_val, sample_processor=None): if data_name not in DATA_REGISTER: available_names = ", ".join(sorted(DATA_REGISTER.keys())) raise ValueError(f"Unknown data {data_name!r}. Available data: {available_names}") - kwargs = {"sample_processor": sample_processor} if data_name == "image_dataset" else {} + kwargs = ( + { + "sample_processor": sample_processor, + "use_training_cache": use_training_cache, + "unconditional_prompt": getattr(sample_processor, "unconditional_prompt", " "), + "expected_cache_info": training_cache_info(config) if use_training_cache else None, + } + if data_name == "image_dataset" + else {} + ) return DATA_REGISTER[data_name](data_config_split, train_or_val=train_or_val, **kwargs) diff --git a/lightx2v_train/train.py b/lightx2v_train/train.py index ef036c521..e201c8dea 100644 --- a/lightx2v_train/train.py +++ b/lightx2v_train/train.py @@ -25,11 +25,17 @@ def main(): try: prepare_data(config) - sample_processor = build_sample_processor(config) + use_training_cache = config["data"].get("use_training_cache", False) + sample_processor = None if use_training_cache else build_sample_processor(config) + dataloader_train = build_data(config, train_or_val="train", sample_processor=sample_processor) + model = build_model(config) - model.load_components() + model.load_components( + load_transformer=True, + load_vae=not use_training_cache, + load_condition_encoder=not use_training_cache, + ) - dataloader_train = build_data(config, train_or_val="train", sample_processor=sample_processor) dataloader_eval = None if config.get("inference", {}).get("infer_every_iters", None): dataloader_eval = build_data(config, train_or_val="val", sample_processor=sample_processor) From 96f08a638dd3f9c792b76e83af53f29d3d894e94 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 16:35:54 +0000 Subject: [PATCH 23/34] update --- lightx2v_train/cache_data.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightx2v_train/cache_data.py b/lightx2v_train/cache_data.py index a0720f21a..e4f0605ab 100644 --- a/lightx2v_train/cache_data.py +++ b/lightx2v_train/cache_data.py @@ -1,11 +1,12 @@ import argparse import torch +from loguru import logger + from lightx2v_train.data import build_data, build_sample_processor, prepare_data 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 -from loguru import logger def parse_args(): From e167ba6904e51ca4f05aca52f135cb30b63d4961 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 16:52:09 +0000 Subject: [PATCH 24/34] remove fastwam --- lightx2v_train/cache_data.py | 10 +- .../libero_plus_uncond_2cam224_4gpu.yaml | 66 -- .../libero_plus_uncond_2cam224_8gpu.yaml | 66 -- .../train/fastwam/libero_uncond_2cam224.yaml | 63 -- .../lightx2v_train/data/__init__.py | 10 - .../lightx2v_train/data/libero/__init__.py | 3 - .../lightx2v_train/data/libero/dataset.py | 222 ----- .../data/libero/lerobot_dataset.py | 187 ---- .../lightx2v_train/data/libero/normalizer.py | 69 -- .../lightx2v_train/data/libero/preparation.py | 307 ------- .../lightx2v_train/data/libero/processor.py | 111 --- .../data/libero/robot_video_dataset.py | 125 --- .../data/libero/video_decoder.py | 67 -- .../lightx2v_train/data/preparation.py | 7 - .../model_capabilities/__init__.py | 2 - .../model_capabilities/training.py | 30 - .../lightx2v_train/model_zoo/__init__.py | 2 - .../model_zoo/native/wan/fastwam/__init__.py | 3 - .../native/wan/fastwam/action_dit.py | 404 --------- .../model_zoo/native/wan/fastwam/layers.py | 100 --- .../model_zoo/native/wan/fastwam/loader.py | 125 --- .../model_zoo/native/wan/fastwam/model.py | 809 ------------------ .../model_zoo/native/wan/fastwam/mot.py | 516 ----------- .../model_zoo/native/wan/fastwam/video_dit.py | 221 ----- .../wan_world_action_capability.py | 40 - .../model_zoo/wan/wan_fastwam.py | 129 --- .../schedulers/flow_matching.py | 64 -- .../lightx2v_train/trainers/__init__.py | 1 - .../lightx2v_train/trainers/world_action.py | 689 --------------- .../lightx2v_train/utils/registry.py | 4 - .../run_fastwam_libero_plus_train_8gpu.sh | 10 - .../scripts/run_fastwam_libero_train.sh | 10 - lightx2v_train/train.py | 10 +- 33 files changed, 14 insertions(+), 4468 deletions(-) delete mode 100644 lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml delete mode 100644 lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml delete mode 100644 lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml delete mode 100644 lightx2v_train/lightx2v_train/data/libero/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/dataset.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/normalizer.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/preparation.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/processor.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py delete mode 100644 lightx2v_train/lightx2v_train/data/libero/video_decoder.py delete mode 100644 lightx2v_train/lightx2v_train/data/preparation.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py delete mode 100644 lightx2v_train/lightx2v_train/trainers/world_action.py delete mode 100755 lightx2v_train/scripts/run_fastwam_libero_plus_train_8gpu.sh delete mode 100755 lightx2v_train/scripts/run_fastwam_libero_train.sh diff --git a/lightx2v_train/cache_data.py b/lightx2v_train/cache_data.py index e4f0605ab..49a658072 100644 --- a/lightx2v_train/cache_data.py +++ b/lightx2v_train/cache_data.py @@ -3,9 +3,14 @@ import torch from loguru import logger -from lightx2v_train.data import build_data, build_sample_processor, prepare_data +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.runtime import ( + cleanup_distributed, + init_distributed, + load_config, + setup_logger, +) from lightx2v_train.trainers import build_trainer @@ -39,7 +44,6 @@ def main(): setup_logger(config) try: - prepare_data(config) sample_processor = build_sample_processor(config) dataloader = build_data(config, train_or_val=args.split, sample_processor=sample_processor) 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 44b094c23..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_4gpu.yaml +++ /dev/null @@ -1,66 +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 - num_workers: 8 - val: - <<: *libero_plus_data - -training: - method: world_action - 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: 32 - 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 fd0db2675..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml +++ /dev/null @@ -1,66 +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 - num_workers: 8 - val: - <<: *libero_plus_data - -training: - method: world_action - 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: 16 - 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 24bed494d..000000000 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ /dev/null @@ -1,63 +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 - num_workers: 8 - val: - <<: *libero_data - -training: - method: world_action - output_dir: runs/fastwam_libero_full - # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule. - max_train_iters: 21700 - gradient_accumulation_iters: 32 - 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/lightx2v_train/data/__init__.py b/lightx2v_train/lightx2v_train/data/__init__.py index f7190c545..d2c344d28 100644 --- a/lightx2v_train/lightx2v_train/data/__init__.py +++ b/lightx2v_train/lightx2v_train/data/__init__.py @@ -6,10 +6,6 @@ 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 - - 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 @@ -18,10 +14,6 @@ def __getattr__(name): "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}") @@ -29,9 +21,7 @@ def __getattr__(name): "build_data", "build_sample_processor", "build_image_dataset", - "build_libero_fastwam_dataset", "build_latent_dataset", "build_prompt_dataset", "build_video_dataset", - "prepare_data", ] 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 a6549b4c4..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=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 f0db2db98..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() - encoding_chunk_size = 16 - with torch.no_grad(): - for start in range(0, len(missing_prompts), encoding_chunk_size): - batch_prompts = missing_prompts[start : start + encoding_chunk_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/model_capabilities/__init__.py b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py index 50fd08241..ada78029b 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/__init__.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/__init__.py @@ -24,7 +24,6 @@ TeacherForcingStepContext, TrainableModelCapability, TrainingCachePayload, - WorldActionTrainingCapability, ) __all__ = [ @@ -51,5 +50,4 @@ "TeacherForcingStepContext", "TrainingCachePayload", "TrainableModelCapability", - "WorldActionTrainingCapability", ] diff --git a/lightx2v_train/lightx2v_train/model_capabilities/training.py b/lightx2v_train/lightx2v_train/model_capabilities/training.py index ea5fe9da2..82a8bfaa1 100644 --- a/lightx2v_train/lightx2v_train/model_capabilities/training.py +++ b/lightx2v_train/lightx2v_train/model_capabilities/training.py @@ -450,33 +450,3 @@ def save_consolidated(self, output_path) -> None: @abstractmethod def save_full_model(self, output_path) -> None: pass - - -class WorldActionTrainingCapability(ModelCapability): - @abstractmethod - def configure(self) -> None: - pass - - @abstractmethod - def parameters(self): - pass - - @abstractmethod - def module(self): - pass - - @abstractmethod - def compute_loss(self, batch: Mapping[str, Any], module=None) -> LossResult: - pass - - @abstractmethod - def evaluation_loss(self, batch: Mapping[str, Any]) -> LossResult: - pass - - @abstractmethod - def load_checkpoint(self, path) -> None: - pass - - @abstractmethod - def save_checkpoint(self, path, step=None) -> None: - pass diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 76804bb5a..faf1ee464 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -14,7 +14,6 @@ "MiniMaxH3T2AVModel": (".minimax_h3.minimax_h3_t2av", "MiniMaxH3T2AVModel"), "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), "QwenImageEditModel": (".qwen_image.qwen_image_edit", "QwenImageEditModel"), - "WanFastWAMModel": (".wan.wan_fastwam", "WanFastWAMModel"), "WanT2VModel": (".wan.wan_t2v", "WanT2VModel"), "WanTI2V5BModel": (".wan.wan_ti2v_5b", "WanTI2V5BModel"), } @@ -64,5 +63,4 @@ def __getattr__(name): "MiniMaxH3T2AVModel", "WanT2VModel", "WanTI2V5BModel", - "WanFastWAMModel", ] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py deleted file mode 100644 index 080d4c67e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .model import FastWAM - -__all__ = ["FastWAM"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py deleted file mode 100644 index b4b317564..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py +++ /dev/null @@ -1,404 +0,0 @@ -import logging -import os -from typing import Any, Dict, Optional - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from ..modules.model2_2 import rope_params -from .layers import expert_block_forward, timestep_embedding - -logger = logging.getLogger(__name__) - - -def _interpolate_last_dim(tensor: torch.Tensor, new_size: int) -> torch.Tensor: - if tensor.shape[-1] == new_size: - return tensor - flat = tensor.reshape(-1, 1, tensor.shape[-1]).to(torch.float32) - flat = F.interpolate(flat, size=new_size, mode="linear", align_corners=True) - return flat.reshape(*tensor.shape[:-1], new_size) - - -def _resize_tensor_to_shape(source: torch.Tensor, target_shape: tuple[int, ...]) -> torch.Tensor: - if tuple(source.shape) == target_shape: - return source - - output = source.to(torch.float32) - while output.ndim < len(target_shape): - output = output.unsqueeze(0) - while output.ndim > len(target_shape): - if output.shape[0] != 1: - raise ValueError(f"Cannot reduce tensor rank for resize: source={tuple(source.shape)} target={target_shape}") - output = output.squeeze(0) - - for dimension, new_size in enumerate(target_shape): - if output.shape[dimension] == new_size: - continue - permutation = [index for index in range(output.ndim) if index != dimension] + [dimension] - inverse_permutation = [0] * output.ndim - for index, original_dimension in enumerate(permutation): - inverse_permutation[original_dimension] = index - permuted = output.permute(*permutation).contiguous() - permuted = _interpolate_last_dim(permuted, new_size) - output = permuted.permute(*inverse_permutation).contiguous() - - if tuple(output.shape) != target_shape: - raise ValueError(f"Resize produced wrong shape: source={tuple(source.shape)} target={target_shape} got={tuple(output.shape)}") - return output.to(dtype=source.dtype) - - -class RMSNorm(nn.Module): - def __init__(self, dim, eps=1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) - - def forward(self, value): - dtype = value.dtype - value = value.float() - value = value * torch.rsqrt(value.pow(2).mean(dim=-1, keepdim=True) + self.eps) - return value.to(dtype) * self.weight - - -class WideAttention(nn.Module): - """Attention projection whose width may differ from the expert hidden width.""" - - def __init__(self, hidden_dim, num_heads, head_dim, eps): - super().__init__() - attention_dim = num_heads * head_dim - self.num_heads = num_heads - self.head_dim = head_dim - self.q = nn.Linear(hidden_dim, attention_dim) - self.k = nn.Linear(hidden_dim, attention_dim) - self.v = nn.Linear(hidden_dim, attention_dim) - self.o = nn.Linear(attention_dim, hidden_dim) - self.norm_q = RMSNorm(attention_dim, eps) - self.norm_k = RMSNorm(attention_dim, eps) - - -class ActionBlock(nn.Module): - def __init__(self, hidden_dim, ffn_dim, num_heads, head_dim, eps): - super().__init__() - self.num_heads = num_heads - self.self_attn = WideAttention(hidden_dim, num_heads, head_dim, eps) - self.cross_attn = WideAttention(hidden_dim, num_heads, head_dim, eps) - self.norm1 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) - self.norm2 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) - self.norm3 = nn.LayerNorm(hidden_dim, eps=eps) - self.ffn = nn.Sequential( - nn.Linear(hidden_dim, ffn_dim), - nn.GELU(approximate="tanh"), - nn.Linear(ffn_dim, hidden_dim), - ) - self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5) - - -class ActionDiT(nn.Module): - ACTION_BACKBONE_SKIP_PREFIXES = ("action_encoder.", "head.") - ACTION_BACKBONE_META_KEYS = ( - "hidden_dim", - "ffn_dim", - "num_layers", - "num_heads", - "attn_head_dim", - "text_dim", - "freq_dim", - "eps", - ) - - def __init__( - self, - hidden_dim: int, - action_dim: int, - ffn_dim: int, - text_dim: int, - freq_dim: int, - eps: float, - num_heads: int, - attn_head_dim: int, - num_layers: int, - use_gradient_checkpointing: bool = False, - ): - super().__init__() - self.hidden_dim = hidden_dim - self.action_dim = action_dim - self.ffn_dim = ffn_dim - self.text_dim = text_dim - self.freq_dim = freq_dim - self.num_heads = num_heads - self.attn_head_dim = attn_head_dim - - if num_heads <= 0: - raise ValueError(f"`num_heads` must be > 0, got {num_heads}") - if attn_head_dim <= 0: - raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") - if attn_head_dim % 2 != 0: - raise ValueError(f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}") - - self.action_encoder = nn.Linear(action_dim, hidden_dim) - self.text_embedding = nn.Sequential( - nn.Linear(text_dim, hidden_dim), - nn.GELU(approximate="tanh"), - nn.Linear(hidden_dim, hidden_dim), - ) - self.time_embedding = nn.Sequential( - nn.Linear(freq_dim, hidden_dim), - nn.SiLU(), - nn.Linear(hidden_dim, hidden_dim), - ) - self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) - self.blocks = nn.ModuleList( - [ - ActionBlock( - hidden_dim=hidden_dim, - ffn_dim=ffn_dim, - num_heads=num_heads, - head_dim=attn_head_dim, - eps=eps, - ) - for _ in range(num_layers) - ] - ) - self.head = nn.Linear(hidden_dim, action_dim) - self.freqs = rope_params(1024, attn_head_dim) - - self.use_gradient_checkpointing = use_gradient_checkpointing - - @classmethod - def backbone_key_set(cls, keys) -> set[str]: - return {key for key in keys if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES)} - - @classmethod - def from_video_expert( - cls, - action_dit_config: dict[str, Any], - video_expert: nn.Module, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - apply_alpha_scaling: bool = True, - ) -> "ActionDiT": - action_expert = cls(**action_dit_config).to(device=device, dtype=torch_dtype) - action_state = action_expert.state_dict() - video_state = video_expert.state_dict() - copied = 0 - interpolated = 0 - - with torch.no_grad(): - for key in sorted(cls.backbone_key_set(action_state.keys())): - if key not in video_state: - raise ValueError(f"ActionDiT backbone key `{key}` does not exist in the Wan video expert.") - source = video_state[key] - target = action_state[key] - if tuple(source.shape) == tuple(target.shape): - value = source - copied += 1 - else: - value = _resize_tensor_to_shape(source, tuple(target.shape)) - if apply_alpha_scaling and source.ndim >= 2 and source.shape[-1] != target.shape[-1]: - alpha = (float(source.shape[-1]) / float(target.shape[-1])) ** 0.5 - value = value.to(torch.float32).mul_(alpha) - interpolated += 1 - target.copy_(value.to(device=target.device, dtype=target.dtype)) - - logger.info( - "Initialized ActionDiT backbone from Wan video expert (copied=%d, interpolated=%d, random_kept_prefixes=%s).", - copied, - interpolated, - list(cls.ACTION_BACKBONE_SKIP_PREFIXES), - ) - return action_expert - - @classmethod - def from_pretrained( - cls, - action_dit_config: dict[str, Any], - action_dit_pretrained_path: str | None = None, - skip_dit_load_from_pretrain: bool = False, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - ) -> "ActionDiT": - if action_dit_config is None: - raise ValueError("`action_dit_config` is required for ActionDiT.from_pretrained().") - if skip_dit_load_from_pretrain: - logger.info("Skipping ActionDiT pretrained load (`skip_dit_load_from_pretrain=True`); initializing action expert randomly and expecting checkpoint override.") - return cls(**action_dit_config).to(device=device, dtype=torch_dtype) - if not action_dit_pretrained_path: - logger.info("No `action_dit_pretrained_path` provided, initializing ActionDiT with random weights.") - return cls(**action_dit_config).to(device=device, dtype=torch_dtype) - from pathlib import Path - - p = Path(action_dit_pretrained_path) - if not p.is_absolute(): - p = Path(__file__).resolve().parents[4] / p - action_dit_pretrained_path = str(p) - if not os.path.isfile(action_dit_pretrained_path): - raise FileNotFoundError(f"`action_dit_pretrained_path` does not exist: {action_dit_pretrained_path}") - - action_cfg = dict(action_dit_config) - action_expert = cls(**action_cfg).to(device=device, dtype=torch_dtype) - action_state = action_expert.state_dict() - expected_backbone_keys = cls.backbone_key_set(action_state.keys()) - - payload = torch.load(action_dit_pretrained_path, map_location="cpu", weights_only=True) - if not isinstance(payload, dict): - raise ValueError(f"Invalid action backbone payload type from {action_dit_pretrained_path}: {type(payload)}") - - policy = payload.get("policy", {}) - if policy: - logger.info(f"ActionDiT backbone payload policy: {policy}") - - meta = payload.get("meta") - expected_meta = { - "hidden_dim": int(action_cfg["hidden_dim"]), - "ffn_dim": int(action_cfg["ffn_dim"]), - "num_layers": int(action_cfg["num_layers"]), - "num_heads": int(action_cfg["num_heads"]), - "attn_head_dim": int(action_cfg["attn_head_dim"]), - "text_dim": int(action_cfg["text_dim"]), - "freq_dim": int(action_cfg["freq_dim"]), - "eps": float(action_cfg["eps"]), - } - for key in cls.ACTION_BACKBONE_META_KEYS: - if key not in meta: - raise ValueError(f"`meta.{key}` missing in {action_dit_pretrained_path}") - expected_value = expected_meta[key] - got_value = meta[key] - if key == "eps": - if abs(float(got_value) - float(expected_value)) > 1e-12: - raise ValueError(f"`meta.{key}` mismatch in {action_dit_pretrained_path}: expected {expected_value}, got {got_value}") - elif int(got_value) != int(expected_value): - raise ValueError(f"`meta.{key}` mismatch in {action_dit_pretrained_path}: expected {expected_value}, got {got_value}") - - backbone_state_dict = payload.get("backbone_state_dict") - if not isinstance(backbone_state_dict, dict): - raise ValueError(f"`backbone_state_dict` must be a dict in {action_dit_pretrained_path}, got {type(backbone_state_dict)}") - - provided_keys = set(backbone_state_dict.keys()) - missing_keys = sorted(expected_backbone_keys - provided_keys) - unexpected_keys = sorted(provided_keys - expected_backbone_keys) - if missing_keys or unexpected_keys: - raise ValueError( - "Action backbone key mismatch in preprocessed payload. " - f"missing={missing_keys[:10]}{'...' if len(missing_keys) > 10 else ''}, " - f"unexpected={unexpected_keys[:10]}{'...' if len(unexpected_keys) > 10 else ''}" - ) - - merged_state = dict(action_state) - for key in expected_backbone_keys: - value = backbone_state_dict[key] - if not isinstance(value, torch.Tensor): - raise ValueError(f"`backbone_state_dict[{key}]` must be torch.Tensor in {action_dit_pretrained_path}, got {type(value)}") - target = merged_state[key] - if tuple(value.shape) != tuple(target.shape): - raise ValueError(f"Shape mismatch for `{key}` in {action_dit_pretrained_path}: expected {tuple(target.shape)}, got {tuple(value.shape)}") - merged_state[key] = value.to(device=target.device, dtype=target.dtype) - - action_expert.load_state_dict(merged_state, strict=True) - logger.info( - "Loaded ActionDiT backbone from %s (keys=%d; random_kept_prefixes=%s).", - action_dit_pretrained_path, - len(expected_backbone_keys), - list(cls.ACTION_BACKBONE_SKIP_PREFIXES), - ) - return action_expert.to(device=device, dtype=torch_dtype) - - def pre_dit( - self, - action_tokens: torch.Tensor, - timestep: torch.Tensor, - context: torch.Tensor, - context_mask: Optional[torch.Tensor] = None, - ) -> Dict[str, Any]: - if action_tokens.ndim != 3: - raise ValueError(f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}") - if action_tokens.shape[2] != self.action_dim: - raise ValueError(f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}") - if timestep.ndim != 1: - raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") - if context.ndim != 3: - raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") - - batch_size = action_tokens.shape[0] - if self.training and batch_size != 1: - raise ValueError("FastWAM action training only supports physical batch size 1.") - if context.shape[0] != batch_size: - raise ValueError(f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}") - if timestep.shape[0] not in (1, batch_size): - raise ValueError(f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}") - if timestep.shape[0] == 1 and batch_size > 1: - timestep = timestep.expand(batch_size) - - if context_mask is None: - context_mask = torch.ones((batch_size, context.shape[1]), dtype=torch.bool, device=context.device) - else: - if context_mask.ndim != 2: - raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") - if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: - raise ValueError(f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}") - - seq_len = action_tokens.shape[1] - if seq_len > self.freqs.shape[0]: - raise ValueError(f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}.") - - t = self.time_embedding(timestep_embedding(self.freq_dim, timestep)) - t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) - - tokens = self.action_encoder(action_tokens) - context_emb = self.text_embedding(context) - context_attn_mask = context_mask.unsqueeze(1).expand(-1, seq_len, -1) - freqs = self.freqs[:seq_len].view(seq_len, 1, -1).to(tokens.device) - - return { - "tokens": tokens, - "freqs": freqs, - "t": t, - "t_mod": t_mod, - "context": context_emb, - "context_mask": context_attn_mask, - "meta": { - "batch_size": batch_size, - "seq_len": seq_len, - }, - } - - def post_dit(self, tokens: torch.Tensor, pre_state: Dict[str, Any]) -> torch.Tensor: - return self.head(tokens) - - def forward( - self, - action_tokens: torch.Tensor, - timestep: torch.Tensor, - context: torch.Tensor, - context_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - pre_state = self.pre_dit( - action_tokens=action_tokens, - timestep=timestep, - context=context, - context_mask=context_mask, - ) - x = pre_state["tokens"] - context = pre_state["context"] - t_mod = pre_state["t_mod"] - freqs = pre_state["freqs"] - context_mask = pre_state["context_mask"] - - for block in self.blocks: - - def run(value, current_block=block): - return expert_block_forward( - current_block, - value, - context, - t_mod, - freqs, - context_mask, - ) - - if self.use_gradient_checkpointing and self.training: - x = torch.utils.checkpoint.checkpoint(run, x, use_reentrant=False) - else: - x = run(x) - - return self.post_dit(x, pre_state) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py deleted file mode 100644 index 6309546bb..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py +++ /dev/null @@ -1,100 +0,0 @@ -import torch -import torch.nn.functional as F - - -def timestep_embedding(dim: int, position: torch.Tensor) -> torch.Tensor: - half = dim // 2 - frequencies = torch.pow( - 10000, - -torch.arange(half, dtype=torch.float64, device=position.device) / half, - ) - embedding = torch.outer(position.to(torch.float64), frequencies) - return torch.cat([torch.cos(embedding), torch.sin(embedding)], dim=1).to(position.dtype) - - -def apply_rope(x: torch.Tensor, frequencies: torch.Tensor) -> torch.Tensor: - dtype = x.dtype - complex_x = torch.view_as_complex(x.to(torch.float64).reshape(*x.shape[:-1], -1, 2)) - complex_frequencies = frequencies.to(device=x.device, dtype=torch.complex128) - return torch.view_as_real(complex_x * complex_frequencies).flatten(-2).to(dtype) - - -def scaled_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - mask: torch.Tensor | None = None, -) -> torch.Tensor: - attention_mask = None - if mask is not None: - attention_mask = mask.to(device=query.device, dtype=torch.bool) - if attention_mask.ndim == 2: - attention_mask = attention_mask.unsqueeze(0).unsqueeze(0) - elif attention_mask.ndim == 3: - attention_mask = attention_mask.unsqueeze(1) - elif attention_mask.ndim != 4: - raise ValueError(f"Attention mask must be 2D, 3D, or 4D, got {attention_mask.ndim}D") - - output = F.scaled_dot_product_attention( - query.transpose(1, 2), - key.transpose(1, 2), - value.transpose(1, 2), - attn_mask=attention_mask, - ) - return output.transpose(1, 2).contiguous() - - -def split_modulation(block, time_modulation: torch.Tensor): - has_sequence_dimension = time_modulation.ndim == 4 - chunk_dimension = 2 if has_sequence_dimension else 1 - values = (block.modulation.to(time_modulation) + time_modulation).chunk(6, dim=chunk_dimension) - if has_sequence_dimension: - values = tuple(value.squeeze(2) for value in values) - return values - - -def layer_norm(module, value: torch.Tensor) -> torch.Tensor: - if not isinstance(module, torch.nn.LayerNorm): - return module(value) - weight = None if module.weight is None else module.weight.float() - bias = None if module.bias is None else module.bias.float() - return F.layer_norm(value.float(), module.normalized_shape, weight, bias, module.eps).to(value.dtype) - - -def masked_cross_attention(block, x, context, context_mask=None): - batch_size, _, num_heads = x.shape[0], x.shape[1], block.cross_attn.num_heads - head_dim = block.cross_attn.head_dim - query = block.cross_attn.norm_q(block.cross_attn.q(x)).view(batch_size, -1, num_heads, head_dim) - key = block.cross_attn.norm_k(block.cross_attn.k(context)).view(batch_size, -1, num_heads, head_dim) - value = block.cross_attn.v(context).view(batch_size, -1, num_heads, head_dim) - output = scaled_attention(query, key, value, context_mask) - return block.cross_attn.o(output.flatten(2)) - - -def expert_block_forward( - block, - x: torch.Tensor, - context: torch.Tensor, - time_modulation: torch.Tensor, - frequencies: torch.Tensor, - context_mask: torch.Tensor | None = None, - self_attention_mask: torch.Tensor | None = None, -) -> torch.Tensor: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = split_modulation(block, time_modulation) - attention_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa - batch_size, sequence_length = attention_input.shape[:2] - num_heads = block.self_attn.num_heads - head_dim = block.self_attn.head_dim - query = block.self_attn.norm_q(block.self_attn.q(attention_input)).view(batch_size, sequence_length, num_heads, head_dim) - key = block.self_attn.norm_k(block.self_attn.k(attention_input)).view(batch_size, sequence_length, num_heads, head_dim) - value = block.self_attn.v(attention_input).view(batch_size, sequence_length, num_heads, head_dim) - attention = scaled_attention( - apply_rope(query, frequencies), - apply_rope(key, frequencies), - value, - self_attention_mask, - ) - x = x + gate_msa * block.self_attn.o(attention.flatten(2)) - x = x + masked_cross_attention(block, layer_norm(block.norm3, x), context, context_mask) - mlp_input = layer_norm(block.norm2, x) * (1 + scale_mlp) + shift_mlp - return x + gate_mlp * block.ffn(mlp_input) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py deleted file mode 100644 index e7048066d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py +++ /dev/null @@ -1,125 +0,0 @@ -import inspect -import logging -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch -from safetensors.torch import load_file - -from ..modules.t5 import T5EncoderModel -from ..modules.vae2_2 import Wan2_2_VAE -from .video_dit import FastWAMVideoDiT - -logger = logging.getLogger(__name__) - - -@dataclass -class Wan22LoadedComponents: - dit: FastWAMVideoDiT - vae: Wan2_2_VAE - text_encoder: T5EncoderModel | None - tokenizer: Any | None - - -def _validate_dit_config(dit_config: dict[str, Any]) -> dict[str, Any]: - if not isinstance(dit_config, dict): - raise ValueError(f"video_dit_config must be a dict, got {type(dit_config)}") - signature = inspect.signature(FastWAMVideoDiT.__init__) - allowed = {name for name in signature.parameters if name != "self"} - required = {name for name, parameter in signature.parameters.items() if name != "self" and parameter.default is inspect.Signature.empty} - unknown = sorted(set(dit_config) - allowed) - missing = sorted(required - set(dit_config)) - if unknown: - raise ValueError(f"Unknown video_dit_config keys: {unknown}") - if missing: - raise ValueError(f"Missing video_dit_config keys: {missing}") - return dict(dit_config) - - -def _require_path(path: Path, description: str, directory: bool = False) -> str: - exists = path.is_dir() if directory else path.is_file() - if not exists: - kind = "directory" if directory else "file" - raise FileNotFoundError(f"Wan model {description} {kind} does not exist: {path}") - return str(path) - - -def _resolve_model_paths(model_path: str): - root = Path(model_path).expanduser().resolve() - _require_path(root, "root", directory=True) - single_dit = root / "diffusion_pytorch_model.safetensors" - if single_dit.is_file(): - dit_path: str | list[str] = str(single_dit) - else: - shards = sorted(root.glob("diffusion_pytorch_model-*-of-*.safetensors")) - if not shards: - raise FileNotFoundError(f"Wan diffusion weights do not exist under: {root}") - dit_path = [str(path) for path in shards] - return ( - dit_path, - _require_path(root / "models_t5_umt5-xxl-enc-bf16.pth", "text encoder"), - _require_path(root / "Wan2.2_VAE.pth", "VAE"), - _require_path(root / "google" / "umt5-xxl", "tokenizer", directory=True), - ) - - -def _load_video_expert(path, config, torch_dtype, device): - model = FastWAMVideoDiT(**config) - paths = path if isinstance(path, list) else [path] - state_dict = {} - for shard_path in paths: - shard = load_file(shard_path, device="cpu") - state_dict.update({key: value.to(torch_dtype) for key, value in shard.items()}) - model.load_state_dict(state_dict, strict=True) - return model.to(device=device, dtype=torch_dtype) - - -def load_wan22_ti2v_5b_components( - model_path: str, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - tokenizer_max_len: int = 512, - dit_config: dict[str, Any] | None = None, - skip_dit_load_from_pretrain: bool = False, - load_text_encoder: bool = True, -): - start = time.time() - if dit_config is None: - raise ValueError("video_dit_config is required for Wan2.2-TI2V-5B loading.") - config = _validate_dit_config(dit_config) - dit_path, text_encoder_path, vae_path, tokenizer_path = _resolve_model_paths(model_path) - - if skip_dit_load_from_pretrain: - dit = FastWAMVideoDiT(**config).to(device=device, dtype=torch_dtype) - else: - dit = _load_video_expert(dit_path, config, torch_dtype, device) - - text_encoder = None - tokenizer = None - if load_text_encoder: - text_encoder = T5EncoderModel( - text_len=int(tokenizer_max_len), - dtype=torch_dtype, - device=torch.device(device), - checkpoint_path=text_encoder_path, - tokenizer_path=tokenizer_path, - ) - tokenizer = text_encoder.tokenizer - - vae = Wan2_2_VAE( - vae_pth=vae_path, - dtype=torch_dtype, - device=device, - ) - vae.model.to(device=device, dtype=torch_dtype) - vae.temporal_downsample_factor = 4 - vae.upsampling_factor = 16 - logger.info("Loaded Wan2.2 FastWAM components in %.2f seconds", time.time() - start) - return Wan22LoadedComponents( - dit=dit, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py deleted file mode 100644 index 023665eb3..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py +++ /dev/null @@ -1,809 +0,0 @@ -import logging -from typing import Any, Optional, Sequence, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -from PIL import Image - -from lightx2v_train.schedulers.flow_matching import WanContinuousFlowMatchScheduler - -from .action_dit import ActionDiT -from .loader import load_wan22_ti2v_5b_components -from .mot import MoT - -logger = logging.getLogger(__name__) - -WAN22_TI2V_5B_VIDEO_CONFIG = { - "has_image_input": False, - "patch_size": (1, 2, 2), - "in_dim": 48, - "hidden_dim": 3072, - "ffn_dim": 14336, - "freq_dim": 256, - "text_dim": 4096, - "out_dim": 48, - "num_heads": 24, - "attn_head_dim": 128, - "num_layers": 30, - "eps": 1e-6, - "seperated_timestep": True, - "fuse_vae_embedding_in_latents": True, - "video_attention_mask_mode": "first_frame_causal", -} - -FASTWAM_ACTION_CONFIG = { - "action_dim": 7, - "hidden_dim": 1024, - "ffn_dim": 4096, - "num_heads": 24, - "attn_head_dim": 128, - "num_layers": 30, - "text_dim": 4096, - "freq_dim": 256, - "eps": 1e-6, -} - - -class FastWAM(torch.nn.Module): - """MoT world model with video/action experts.""" - - def __init__( - self, - video_expert, - action_expert: ActionDiT, - mot: MoT, - vae, - text_encoder=None, - tokenizer=None, - text_dim: Optional[int] = None, - proprio_dim: Optional[int] = None, - device: str = "cpu", - torch_dtype: torch.dtype = torch.float32, - video_train_shift: float = 5.0, - video_infer_shift: float = 5.0, - video_num_train_timesteps: int = 1000, - action_train_shift: float = 5.0, - action_infer_shift: float = 5.0, - action_num_train_timesteps: int = 1000, - loss_lambda_video: float = 1.0, - loss_lambda_action: float = 1.0, - ): - super().__init__() - self.video_expert = video_expert - self.action_expert = action_expert - self.mot = mot - # Keep trainer compatibility: optimizer and freeze logic use `model.dit`. - self.dit = self.mot - - self.vae = vae - self.text_encoder = text_encoder - self.tokenizer = tokenizer - if text_dim is None: - if self.text_encoder is None: - raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") - text_dim = int(self.text_encoder.model.dim) - self.text_dim = int(text_dim) - self.proprio_dim = None if proprio_dim is None else int(proprio_dim) - if self.proprio_dim is not None: - self.proprio_encoder = nn.Linear(self.proprio_dim, self.text_dim).to(torch_dtype) - else: - self.proprio_encoder = None - - self.train_video_scheduler = WanContinuousFlowMatchScheduler( - num_train_timesteps=video_num_train_timesteps, - shift=video_train_shift, - ) - self.infer_video_scheduler = WanContinuousFlowMatchScheduler( - num_train_timesteps=video_num_train_timesteps, - shift=video_infer_shift, - ) - self.train_action_scheduler = WanContinuousFlowMatchScheduler( - num_train_timesteps=action_num_train_timesteps, - shift=action_train_shift, - ) - self.infer_action_scheduler = WanContinuousFlowMatchScheduler( - num_train_timesteps=action_num_train_timesteps, - shift=action_infer_shift, - ) - self.device = torch.device(device) - self.torch_dtype = torch_dtype - self.loss_lambda_video = float(loss_lambda_video) - self.loss_lambda_action = float(loss_lambda_action) - - self.to(self.device) - - @classmethod - def from_wan22_pretrained( - cls, - model_path: str, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - tokenizer_max_len: int = 512, - load_text_encoder: bool = True, - proprio_dim: Optional[int] = None, - video_dit_config: dict[str, Any] | None = None, - action_dit_config: dict[str, Any] | None = None, - action_dit_pretrained_path: str | None = None, - skip_dit_load_from_pretrain: bool = False, - mot_checkpoint_mixed_attn: bool = True, - video_train_shift: float = 5.0, - video_infer_shift: float = 5.0, - video_num_train_timesteps: int = 1000, - action_train_shift: float = 5.0, - action_infer_shift: float = 5.0, - action_num_train_timesteps: int = 1000, - loss_lambda_video: float = 1.0, - loss_lambda_action: float = 1.0, - ): - video_dit_config = {**WAN22_TI2V_5B_VIDEO_CONFIG, **(video_dit_config or {})} - action_dit_config = {**FASTWAM_ACTION_CONFIG, **(action_dit_config or {})} - - components = load_wan22_ti2v_5b_components( - model_path=model_path, - device=device, - torch_dtype=torch_dtype, - tokenizer_max_len=tokenizer_max_len, - dit_config=video_dit_config, - skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, - load_text_encoder=load_text_encoder, - ) - - video_expert = components.dit - if action_dit_pretrained_path or skip_dit_load_from_pretrain: - action_expert = ActionDiT.from_pretrained( - action_dit_config=action_dit_config, - action_dit_pretrained_path=action_dit_pretrained_path, - skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, - device=device, - torch_dtype=torch_dtype, - ) - else: - action_expert = ActionDiT.from_video_expert( - action_dit_config=action_dit_config, - video_expert=video_expert, - device=device, - torch_dtype=torch_dtype, - ) - if int(action_expert.num_heads) != int(video_expert.num_heads): - raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") - if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): - raise ValueError("ActionDiT `attn_head_dim` must match video expert for MoT mixed attention.") - if int(len(action_expert.blocks)) != int(len(video_expert.blocks)): - raise ValueError("ActionDiT `num_layers` must match video expert.") - - mot = MoT( - mixtures={"video": video_expert, "action": action_expert}, - mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, - ) - - model = cls( - video_expert=video_expert, - action_expert=action_expert, - mot=mot, - vae=components.vae, - text_encoder=components.text_encoder, - tokenizer=components.tokenizer, - text_dim=int(video_dit_config["text_dim"]), - proprio_dim=proprio_dim, - device=device, - torch_dtype=torch_dtype, - video_train_shift=video_train_shift, - video_infer_shift=video_infer_shift, - video_num_train_timesteps=video_num_train_timesteps, - action_train_shift=action_train_shift, - action_infer_shift=action_infer_shift, - action_num_train_timesteps=action_num_train_timesteps, - loss_lambda_video=loss_lambda_video, - loss_lambda_action=loss_lambda_action, - ) - return model - - def to(self, *args, **kwargs): - super().to(*args, **kwargs) - self.mot.to(*args, **kwargs) - if self.text_encoder is not None: - self.text_encoder.model.to(*args, **kwargs) - self.vae.model.to(*args, **kwargs) - return self - - @staticmethod - def _check_resize_height_width(height, width, num_frames): - if height % 16 != 0: - height = (height + 15) // 16 * 16 - if width % 16 != 0: - width = (width + 15) // 16 * 16 - if num_frames % 4 != 1: - num_frames = (num_frames + 3) // 4 * 4 + 1 - return height, width, num_frames - - @torch.no_grad() - def encode_prompt(self, prompt: Union[str, Sequence[str]]): - if self.text_encoder is None or self.tokenizer is None: - raise ValueError("Prompt encoding requires loaded text encoder/tokenizer. Set `load_text_encoder=true` or provide precomputed `context/context_mask`.") - ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) - ids = ids.to(self.device) - mask = mask.to(self.device, dtype=torch.bool) - prompt_emb = self.text_encoder.model(ids, mask) - # FIXME: original implementation's zero padding is visible in cross-attn. - seq_lens = mask.gt(0).sum(dim=1).long() - for i, v in enumerate(seq_lens): - prompt_emb[i, v:] = 0 - mask = torch.ones_like(mask) - return prompt_emb.to(device=self.device), mask - - def _append_proprio_to_context( - self, - context: torch.Tensor, - context_mask: torch.Tensor, - proprio: Optional[torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.proprio_encoder is None or proprio is None: - return context, context_mask - if proprio.ndim != 2: - raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") - if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: - raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") - proprio_token = self.proprio_encoder(proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1)).to(dtype=context.dtype) # [B, 1, D] - proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) - return ( - torch.cat([context, proprio_token], dim=1), - torch.cat([context_mask, proprio_mask], dim=1), - ) - - @torch.no_grad() - def _encode_video_latents(self, video_tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): - del tile_size, tile_stride - if tiled: - raise ValueError("Wan2.2 training VAE does not support tiled encoding.") - latents = self.vae.encode([video for video in video_tensor]) - return torch.stack(latents).to(device=self.device, dtype=self.torch_dtype) - - @torch.no_grad() - def _encode_input_image_latents_tensor(self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): - if input_image.ndim == 3: - input_image = input_image.unsqueeze(0) - if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: - raise ValueError(f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}") - image = input_image.to(device=self.device)[0].unsqueeze(1) - del tile_size, tile_stride - if tiled: - raise ValueError("Wan2.2 training VAE does not support tiled encoding.") - return torch.stack(self.vae.encode([image])).to(device=self.device, dtype=self.torch_dtype) - - def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): - del tile_size, tile_stride - if tiled: - raise ValueError("Wan2.2 training VAE does not support tiled decoding.") - video_tensor = torch.stack(self.vae.decode([latent for latent in latents])) - video_tensor = video_tensor.squeeze(0).detach().float().clamp(-1, 1) - video_tensor = ((video_tensor + 1.0) * 127.5).to(torch.uint8).cpu() - frames = [] - for t in range(video_tensor.shape[1]): - frame = video_tensor[:, t].permute(1, 2, 0).numpy() - frames.append(Image.fromarray(frame)) - return frames - - def build_inputs(self, sample, tiled: bool = False): - video = sample["video"] - if "context" not in sample or "context_mask" not in sample: - raise ValueError("FastWAM training requires `sample['context']` and `sample['context_mask']`.") - context = sample["context"] - context_mask = sample["context_mask"] - proprio = sample.get("proprio", None) - if video.ndim != 5: - raise ValueError(f"`sample['video']` must be 5D [B, 3, T, H, W], got shape {tuple(video.shape)}") - if video.shape[1] != 3: - raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") - - _, _, num_frames, height, width = video.shape - if video.shape[0] != 1: - raise ValueError("FastWAM training only supports physical batch size 1.") - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"Video spatial dims must be multiples of 16, got H={height}, W={width}") - if num_frames % 4 != 1: - raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") - if num_frames <= 1: - raise ValueError(f"Video T must be > 1 for action-conditioned training, got T={num_frames}") - - if "action" not in sample: - raise ValueError("`sample['action']` is required for FastWAM training.") - - action = sample["action"] - if action.ndim != 3: - raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") - action_horizon = int(action.shape[1]) - if action_horizon % (num_frames - 1) != 0: - raise ValueError(f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}") - - action_is_pad = sample.get("action_is_pad", None) - if action_is_pad is not None: - if action_is_pad.ndim != 2: - raise ValueError(f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}") - if action_is_pad.shape[0] != 1 or action_is_pad.shape[1] != action_horizon: - raise ValueError(f"`sample['action_is_pad']` shape mismatch: got {tuple(action_is_pad.shape)} vs expected (1, {action_horizon})") - - image_is_pad = sample.get("image_is_pad", None) - if image_is_pad is not None: - if image_is_pad.ndim != 2: - raise ValueError(f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}") - if image_is_pad.shape[0] != 1 or image_is_pad.shape[1] != num_frames: - raise ValueError(f"`sample['image_is_pad']` shape mismatch: got {tuple(image_is_pad.shape)} vs expected (1, {num_frames})") - - input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) - input_latents = self._encode_video_latents(input_video, tiled=tiled) - - first_frame_latents = None - fuse_flag = False - if getattr(self.video_expert, "fuse_vae_embedding_in_latents", False): - first_frame_latents = input_latents[:, :, 0:1] - fuse_flag = True - - if context.ndim != 3 or context_mask.ndim != 2: - raise ValueError(f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}") - context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) - context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) - if self.proprio_encoder is not None: - if proprio is None: - raise ValueError("`sample['proprio']` is required when `proprio_dim` is enabled.") - if proprio.ndim != 3: - raise ValueError(f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}") - if proprio.shape[2] != self.proprio_dim: - raise ValueError(f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}") - proprio = proprio[:, 0, :] # [B, D] - context, context_mask = self._append_proprio_to_context( - context=context, - context_mask=context_mask, - proprio=proprio.to(device=self.device, dtype=self.torch_dtype), - ) - action = action.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) - - if action_is_pad is not None: - action_is_pad = action_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) - if image_is_pad is not None: - image_is_pad = image_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) - - return { - "context": context, - "context_mask": context_mask, - "input_latents": input_latents, - "first_frame_latents": first_frame_latents, - "fuse_vae_embedding_in_latents": fuse_flag, - "action": action, - "action_is_pad": action_is_pad, - "image_is_pad": image_is_pad, - } - - @torch.no_grad() - def _build_mot_attention_mask( - self, - video_seq_len: int, - action_seq_len: int, - video_tokens_per_frame: int, - device: torch.device, - ) -> torch.Tensor: - total_seq_len = video_seq_len + action_seq_len - mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) - - # video -> video - mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( - video_seq_len=video_seq_len, - video_tokens_per_frame=video_tokens_per_frame, - device=device, - ) - # action -> action - mask[video_seq_len:, video_seq_len:] = True - # action -> first-frame video only - first_frame_tokens = min(video_tokens_per_frame, video_seq_len) - mask[video_seq_len:, :first_frame_tokens] = True - return mask - - def _compute_video_loss_per_sample( - self, - pred_video: torch.Tensor, - target_video: torch.Tensor, - image_is_pad: Optional[torch.Tensor], - include_initial_video_step: bool, - ) -> torch.Tensor: - video_loss_token = F.mse_loss(pred_video.float(), target_video.float(), reduction="none").mean(dim=(1, 3, 4)) - if image_is_pad is None: - return video_loss_token.mean(dim=1) - - temporal_factor = int(self.vae.temporal_downsample_factor) - if temporal_factor <= 0: - raise ValueError(f"`vae.temporal_downsample_factor` must be positive, got {temporal_factor}.") - if image_is_pad.shape[1] < 1: - raise ValueError("`image_is_pad` must contain at least one frame.") - if (image_is_pad.shape[1] - 1) % temporal_factor != 0: - raise ValueError(f"Cannot align `image_is_pad` with video latent steps: num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}.") - - tail_is_pad = image_is_pad[:, 1:] - latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) - if include_initial_video_step: - video_is_pad = torch.cat([image_is_pad[:, :1], latent_tail_is_pad], dim=1) - else: - video_is_pad = latent_tail_is_pad - - if video_is_pad.shape[1] != video_loss_token.shape[1]: - raise ValueError(f"Video-loss mask shape mismatch: mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}.") - - valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) - valid_sum = valid.sum(dim=1).clamp(min=1.0) - return (video_loss_token * valid).sum(dim=1) / valid_sum - - def training_loss(self, sample, tiled: bool = False): - inputs = self.build_inputs(sample, tiled=tiled) - input_latents = inputs["input_latents"] - if input_latents.shape[0] != 1: - raise ValueError("FastWAM training only supports physical batch size 1.") - context = inputs["context"] - context_mask = inputs["context_mask"] - action = inputs["action"] - action_is_pad = inputs["action_is_pad"] - image_is_pad = inputs["image_is_pad"] - - noise_video = torch.randn_like(input_latents) - timestep_video = self.train_video_scheduler.sample_training_t( - device=self.device, - dtype=input_latents.dtype, - ) - latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video) - target_video = self.train_video_scheduler.training_target(input_latents, noise_video, timestep_video) - - if inputs["first_frame_latents"] is not None: - latents[:, :, 0:1] = inputs["first_frame_latents"] - - noise_action = torch.randn_like(action) - timestep_action = self.train_action_scheduler.sample_training_t( - device=self.device, - dtype=action.dtype, - ) - noisy_action = self.train_action_scheduler.add_noise(action, noise_action, timestep_action) - target_action = self.train_action_scheduler.training_target(action, noise_action, timestep_action) - - video_pre = self.video_expert.pre_dit( - x=latents, - timestep=timestep_video, - context=context, - context_mask=context_mask, - action=action, - fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"], - ) - - action_pre = self.action_expert.pre_dit( - action_tokens=noisy_action, - timestep=timestep_action, - context=context, - context_mask=context_mask, - ) - - video_tokens = video_pre["tokens"] - action_tokens = action_pre["tokens"] - - attention_mask = self._build_mot_attention_mask( - video_seq_len=video_tokens.shape[1], - action_seq_len=action_tokens.shape[1], - video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), - device=video_tokens.device, - ) - tokens_out = self.mot( - embeds_all={ - "video": video_tokens, - "action": action_tokens, - }, - attention_mask=attention_mask, - freqs_all={ - "video": video_pre["freqs"], - "action": action_pre["freqs"], - }, - context_all={ - "video": { - "context": video_pre["context"], - "mask": video_pre["context_mask"], - }, - "action": { - "context": action_pre["context"], - "mask": action_pre["context_mask"], - }, - }, - t_mod_all={ - "video": video_pre["t_mod"], - "action": action_pre["t_mod"], - }, - ) - - pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) - - pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) - - include_initial_video_step = inputs["first_frame_latents"] is None - if inputs["first_frame_latents"] is not None: - pred_video = pred_video[:, :, 1:] - target_video = target_video[:, :, 1:] - - loss_video_per_sample = self._compute_video_loss_per_sample( - pred_video=pred_video, - target_video=target_video, - image_is_pad=image_is_pad, - include_initial_video_step=include_initial_video_step, - ) - video_weight = self.train_video_scheduler.training_weight(timestep_video).to(loss_video_per_sample.device, dtype=loss_video_per_sample.dtype) - loss_video = (loss_video_per_sample * video_weight).mean() - - action_loss_token = F.mse_loss(pred_action.float(), target_action.float(), reduction="none").mean(dim=2) # [B, T] - if action_is_pad is not None: - valid = (~action_is_pad).to(device=action_loss_token.device, dtype=action_loss_token.dtype) - valid_sum = valid.sum(dim=1).clamp(min=1.0) - action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum - else: - action_loss_per_sample = action_loss_token.mean(dim=1) - - action_weight = self.train_action_scheduler.training_weight(timestep_action).to(action_loss_per_sample.device, dtype=action_loss_per_sample.dtype) - loss_action = (action_loss_per_sample * action_weight).mean() - - loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action - loss_dict = { - "loss_video": self.loss_lambda_video * float(loss_video.detach().item()), - "loss_action": self.loss_lambda_action * float(loss_action.detach().item()), - } - return loss_total, loss_dict - - @torch.no_grad() - def _predict_joint_noise( - self, - latents_video: torch.Tensor, - latents_action: torch.Tensor, - timestep_video: torch.Tensor, - timestep_action: torch.Tensor, - context: torch.Tensor, - context_mask: torch.Tensor, - fuse_vae_embedding_in_latents: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - video_pre = self.video_expert.pre_dit( - x=latents_video, - timestep=timestep_video, - context=context, - context_mask=context_mask, - action=None, - fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, - ) - action_pre = self.action_expert.pre_dit( - action_tokens=latents_action, - timestep=timestep_action, - context=context, - context_mask=context_mask, - ) - - attention_mask = self._build_mot_attention_mask( - video_seq_len=video_pre["tokens"].shape[1], - action_seq_len=action_pre["tokens"].shape[1], - video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), - device=video_pre["tokens"].device, - ) - - tokens_out = self.mot( - embeds_all={ - "video": video_pre["tokens"], - "action": action_pre["tokens"], - }, - attention_mask=attention_mask, - freqs_all={ - "video": video_pre["freqs"], - "action": action_pre["freqs"], - }, - context_all={ - "video": { - "context": video_pre["context"], - "mask": video_pre["context_mask"], - }, - "action": { - "context": action_pre["context"], - "mask": action_pre["context_mask"], - }, - }, - t_mod_all={ - "video": video_pre["t_mod"], - "action": action_pre["t_mod"], - }, - ) - - pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) - pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) - return pred_video, pred_action - - @torch.no_grad() - def infer_joint( - self, - prompt: Optional[str], - input_image: torch.Tensor, - num_video_frames: int, - action_horizon: int, - proprio: Optional[torch.Tensor] = None, - context: Optional[torch.Tensor] = None, - context_mask: Optional[torch.Tensor] = None, - num_inference_steps: int = 20, - sigma_shift: Optional[float] = None, - seed: Optional[int] = None, - rand_device: str = "cpu", - tiled: bool = False, - ) -> dict[str, Any]: - self.eval() - if input_image.ndim == 3: - input_image = input_image.unsqueeze(0) - if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: - raise ValueError(f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}") - _, _, height, width = input_image.shape - checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) - if (checked_h, checked_w) != (height, width): - raise ValueError(f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})") - if checked_t != num_video_frames: - raise ValueError(f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}") - if proprio is not None: - if self.proprio_dim is None: - raise ValueError("`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled.") - if proprio.ndim == 1: - proprio = proprio.unsqueeze(0) - elif proprio.ndim == 2 and proprio.shape[0] == 1: - pass - else: - raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") - if proprio.shape[1] != self.proprio_dim: - raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") - proprio = proprio.to(device=self.device, dtype=self.torch_dtype) - - latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 - latent_h = height // self.vae.upsampling_factor - latent_w = width // self.vae.upsampling_factor - - video_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) - action_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) - latents_video = torch.randn( - (1, self.vae.model.z_dim, latent_t, latent_h, latent_w), - generator=video_generator, - device=rand_device, - dtype=torch.float32, - ).to(device=self.device, dtype=self.torch_dtype) - latents_action = torch.randn( - (1, action_horizon, self.action_expert.action_dim), - generator=action_generator, - device=rand_device, - dtype=torch.float32, - ).to(device=self.device, dtype=self.torch_dtype) - - input_image = input_image.to(device=self.device, dtype=self.torch_dtype) - first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) - latents_video[:, :, 0:1] = first_frame_latents.clone() - fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) - - use_prompt = prompt is not None - use_context = context is not None or context_mask is not None - if use_prompt and use_context: - raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") - if not use_prompt and not use_context: - raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") - - if use_prompt: - context, context_mask = self.encode_prompt(prompt) - else: - if context is None or context_mask is None: - raise ValueError("`context` and `context_mask` must be both provided together.") - if context.ndim == 2: - context = context.unsqueeze(0) - if context_mask.ndim == 1: - context_mask = context_mask.unsqueeze(0) - if context.ndim != 3 or context_mask.ndim != 2: - raise ValueError(f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}") - context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) - context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) - if proprio is not None: - context, context_mask = self._append_proprio_to_context( - context=context, - context_mask=context_mask, - proprio=proprio, - ) - - infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( - num_inference_steps=num_inference_steps, - device=self.device, - dtype=latents_video.dtype, - shift_override=sigma_shift, - ) - infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( - num_inference_steps=num_inference_steps, - device=self.device, - dtype=latents_action.dtype, - shift_override=sigma_shift, - ) - for step_t_video, step_delta_video, step_t_action, step_delta_action in zip( - infer_timesteps_video, - infer_deltas_video, - infer_timesteps_action, - infer_deltas_action, - ): - timestep_video = step_t_video.unsqueeze(0).to(dtype=latents_video.dtype, device=self.device) - timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) - - pred_video_posi, pred_action_posi = self._predict_joint_noise( - latents_video=latents_video, - latents_action=latents_action, - timestep_video=timestep_video, - timestep_action=timestep_action, - context=context, - context_mask=context_mask, - fuse_vae_embedding_in_latents=fuse_flag, - ) - pred_video = pred_video_posi - pred_action = pred_action_posi - - latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) - latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) - latents_video[:, :, 0:1] = first_frame_latents.clone() - - action_out = latents_action[0].detach().to(device="cpu", dtype=torch.float32) - return { - "video": self._decode_latents(latents_video, tiled=tiled), - "action": action_out, - } - - @torch.no_grad() - def infer( - self, - prompt: Optional[str], - input_image: torch.Tensor, - num_frames: int, - action_horizon: int, - proprio: Optional[torch.Tensor] = None, - context: Optional[torch.Tensor] = None, - context_mask: Optional[torch.Tensor] = None, - num_inference_steps: int = 20, - sigma_shift: Optional[float] = None, - seed: Optional[int] = None, - rand_device: str = "cpu", - tiled: bool = False, - ): - return self.infer_joint( - prompt=prompt, - input_image=input_image, - num_video_frames=num_frames, - action_horizon=action_horizon, - proprio=proprio, - context=context, - context_mask=context_mask, - num_inference_steps=num_inference_steps, - sigma_shift=sigma_shift, - seed=seed, - rand_device=rand_device, - tiled=tiled, - ) - - def save_checkpoint(self, path, step=None): - payload = { - "mot": self.mot.state_dict(), - "step": step, - "torch_dtype": str(self.torch_dtype), - } - if self.proprio_encoder is not None: - payload["proprio_encoder"] = self.proprio_encoder.state_dict() - torch.save(payload, path) - - def load_checkpoint(self, path): - payload = torch.load(path, map_location="cpu", weights_only=True) - if "mot" in payload: - self.mot.load_state_dict(payload["mot"], strict=False) - elif "dit" in payload: - logger.warning("Loading legacy `dit` checkpoint into video expert only.") - self.video_expert.load_state_dict(payload["dit"], strict=False) - else: - raise ValueError(f"Checkpoint missing both `mot` and `dit` keys: {path}") - if self.proprio_encoder is not None: - if "proprio_encoder" in payload: - self.proprio_encoder.load_state_dict(payload["proprio_encoder"], strict=True) - else: - logger.warning("Checkpoint has no `proprio_encoder` weights; keeping current `proprio_encoder` params.") - elif "proprio_encoder" in payload: - logger.warning("Checkpoint contains `proprio_encoder` weights but current model has `proprio_dim=None`; ignoring.") - - return payload - - def forward(self, *args, **kwargs): - return self.training_loss(*args, **kwargs) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py deleted file mode 100644 index 9d2e8495d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py +++ /dev/null @@ -1,516 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Dict, Optional - -import torch -import torch.nn as nn - -from .layers import apply_rope, layer_norm, masked_cross_attention, scaled_attention, split_modulation - -logger = logging.getLogger(__name__) - - -class MoT(nn.Module): - def __init__( - self, - mixtures: Dict[str, nn.Module], - mot_checkpoint_mixed_attn: bool = True, - ): - super().__init__() - if not mixtures: - raise ValueError("`mixtures` cannot be empty.") - if "video" not in mixtures or "action" not in mixtures: - raise ValueError("`mixtures` must include both 'video' and 'action' experts.") - - self.mixtures = nn.ModuleDict(mixtures) - self.expert_order = list(self.mixtures.keys()) - self.mot_checkpoint_mixed_attn = mot_checkpoint_mixed_attn - if mot_checkpoint_mixed_attn: - logger.info("Using gradient checkpointing for mixture attention. This will save memory but use more computation.") - - first_expert = self.mixtures[self.expert_order[0]] - self.num_layers = len(first_expert.blocks) - self.num_heads = first_expert.num_heads - self.attn_head_dim = first_expert.attn_head_dim - - for name in self.expert_order[1:]: - expert = self.mixtures[name] - if len(expert.blocks) != self.num_layers: - raise ValueError(f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}") - if expert.num_heads != self.num_heads: - raise ValueError(f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}") - if expert.attn_head_dim != self.attn_head_dim: - raise ValueError(f"All experts must have same attn_head_dim; got {self.attn_head_dim} and {expert.attn_head_dim}") - - logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") - for name in self.expert_order: - expert = self.mixtures[name] - logger.info(f" Expert '{name}': num_params={sum(p.numel() for p in expert.parameters()) / 1e9:.2f} B") - - @staticmethod - def _split_modulation(block, t_mod: torch.Tensor): - return split_modulation(block, t_mod) - - def _mixed_attention( - self, - q_cat: torch.Tensor, - k_cat: torch.Tensor, - v_cat: torch.Tensor, - attention_mask: torch.Tensor, - ) -> torch.Tensor: - def _forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - return scaled_attention(q, k, v, attention_mask) - - if self.mot_checkpoint_mixed_attn and self.training: - return torch.utils.checkpoint.checkpoint( - _forward, - q_cat, - k_cat, - v_cat, - use_reentrant=False, - ) - return _forward(q_cat, k_cat, v_cat) - - @staticmethod - def _apply_expert_post_block( - block, - residual_x: torch.Tensor, - mixed_attn_out: torch.Tensor, - gate_msa: torch.Tensor, - shift_mlp: torch.Tensor, - scale_mlp: torch.Tensor, - gate_mlp: torch.Tensor, - context_payload: Optional[dict], - ) -> torch.Tensor: - x = residual_x + gate_msa * block.self_attn.o(mixed_attn_out.flatten(2)) - - if context_payload is not None: - context = context_payload.get("context") - if context is not None: - x = x + masked_cross_attention( - block, - layer_norm(block.norm3, x), - context, - context_payload.get("mask"), - ) - - mlp_input = layer_norm(block.norm2, x) * (1 + scale_mlp) + shift_mlp - return x + gate_mlp * block.ffn(mlp_input) - - def _build_expert_attention_io( - self, - expert, - block, - x: torch.Tensor, - freqs: torch.Tensor, - t_mod: torch.Tensor, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - bool, - ]: - """Build per-expert attention tensors and post-block states. - - Args: - expert: Expert module that owns this `block`; only used to read - `use_gradient_checkpointing`. - block: Transformer block for current layer (`expert.blocks[layer_idx]`). - x: Current expert tokens, shape [B, S, D]. - freqs: RoPE frequencies aligned with token sequence, shape [S, 1, rope_dim]. - t_mod: Time modulation tensor for this expert/layer. - - Returns: - q: Query after q-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. - k: Key after k-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. - v: Value after v-proj, shape [B, S, H*Dh]. - residual_x: Original input `x` for residual path in post block. - gate_msa: Gating tensor for self-attention residual branch. - shift_mlp: Shift tensor for MLP modulation. - scale_mlp: Scale tensor for MLP modulation. - gate_mlp: Gating tensor for MLP residual branch. - use_gradient_checkpointing: Whether this expert enables checkpointing. - """ - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._split_modulation(block, t_mod) - attn_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa - - batch_size, sequence_length = attn_input.shape[:2] - q = block.self_attn.norm_q(block.self_attn.q(attn_input)).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) - k = block.self_attn.norm_k(block.self_attn.k(attn_input)).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) - v = block.self_attn.v(attn_input).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) - - q = apply_rope(q, freqs) - k = apply_rope(k, freqs) - - use_gradient_checkpointing = bool(getattr(expert, "use_gradient_checkpointing", False)) - return ( - q, - k, - v, - x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - use_gradient_checkpointing, - ) - - def _apply_post_with_optional_checkpoint( - self, - block, - residual_x: torch.Tensor, - gate_msa: torch.Tensor, - shift_mlp: torch.Tensor, - scale_mlp: torch.Tensor, - gate_mlp: torch.Tensor, - use_gradient_checkpointing: bool, - mixed_slice: torch.Tensor, - context_payload: Optional[dict], - ) -> torch.Tensor: - """Apply post-attention computations, with optional checkpointing. - - Args: - block: Transformer block for current layer. - residual_x: Residual input tokens before attention update, shape [B, S, D]. - gate_msa: Gating tensor used after mixed self-attention. - shift_mlp: Shift tensor for MLP input modulation. - scale_mlp: Scale tensor for MLP input modulation. - gate_mlp: Gating tensor used after MLP. - use_gradient_checkpointing: If True and training, checkpoint this post block. - mixed_slice: Mixed-attention output for this expert, shape [B, S, H*Dh]. - context_payload: Optional dict for cross-attention. - - `context`: encoder states [B, L, D] - - `mask`: attention mask [B, S, L] or [B, 1, S, L] - - Returns: - Updated expert tokens after self-attn residual, optional cross-attn, and MLP. - """ - - def _post_fn( - _mixed_slice: torch.Tensor, - _x: torch.Tensor, - _gate_msa: torch.Tensor, - _shift_mlp: torch.Tensor, - _scale_mlp: torch.Tensor, - _gate_mlp: torch.Tensor, - _block=block, - _context_payload=context_payload, - ) -> torch.Tensor: - return self._apply_expert_post_block( - block=_block, - residual_x=_x, - mixed_attn_out=_mixed_slice, - gate_msa=_gate_msa, - shift_mlp=_shift_mlp, - scale_mlp=_scale_mlp, - gate_mlp=_gate_mlp, - context_payload=_context_payload, - ) - - if use_gradient_checkpointing and self.training: - return torch.utils.checkpoint.checkpoint( - _post_fn, - mixed_slice, - residual_x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - use_reentrant=False, - ) - return _post_fn( - mixed_slice, - residual_x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - ) - - def prefill_video_cache( - self, - video_tokens: torch.Tensor, - video_freqs: torch.Tensor, - video_t_mod: torch.Tensor, - video_context_payload: Optional[dict], - video_attention_mask: torch.Tensor, - ) -> list[dict[str, torch.Tensor]]: - """Prefill video branch once and cache per-layer K/V for action denoising. - - Args: - video_tokens: Video tokens before layer 0, shape [B, Sv, D]. - video_freqs: Video RoPE frequencies, shape [Sv, 1, rope_dim]. - video_t_mod: Video time modulation tensor. - video_context_payload: Optional dict for video cross-attention. - - `context`: encoder states [B, L, D] - - `mask`: attention mask [B, Sv, L] or [B, 1, Sv, L] - video_attention_mask: Video self-attention mask, shape [Sv, Sv]. - - Returns: - Layer-wise cache list with length `num_layers`. - Each entry contains: - - `k`: video key tensor [B, Sv, H*Dh] - - `v`: video value tensor [B, Sv, H*Dh] - """ - if "video" not in self.mixtures: - raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") - if video_attention_mask.ndim != 2: - raise ValueError(f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}") - if video_attention_mask.shape[0] != video_attention_mask.shape[1]: - raise ValueError(f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}") - if video_attention_mask.shape[0] != video_tokens.shape[1]: - raise ValueError(f"`video_attention_mask` seq length mismatch: mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}") - - expert = self.mixtures["video"] - x = video_tokens - kv_cache: list[dict[str, torch.Tensor]] = [] - for layer_idx in range(self.num_layers): - block = expert.blocks[layer_idx] - # Build video Q/K/V from current layer input tokens. - ( - q, - k, - v, - residual_x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - use_gradient_checkpointing, - ) = self._build_expert_attention_io( - expert=expert, - block=block, - x=x, - freqs=video_freqs, - t_mod=video_t_mod, - ) - # Video prefill uses only video self-attention mask. - mixed = self._mixed_attention( - q_cat=q, - k_cat=k, - v_cat=v, - attention_mask=video_attention_mask, - ) - # Update video tokens for the next layer and persist current layer K/V. - x = self._apply_post_with_optional_checkpoint( - block=block, - residual_x=residual_x, - gate_msa=gate_msa, - shift_mlp=shift_mlp, - scale_mlp=scale_mlp, - gate_mlp=gate_mlp, - use_gradient_checkpointing=use_gradient_checkpointing, - mixed_slice=mixed, - context_payload=video_context_payload, - ) - kv_cache.append({"k": k, "v": v}) - return kv_cache - - def forward_action_with_video_cache( - self, - action_tokens: torch.Tensor, - action_freqs: torch.Tensor, - action_t_mod: torch.Tensor, - action_context_payload: Optional[dict], - video_kv_cache: list[dict[str, torch.Tensor]], - attention_mask: torch.Tensor, - video_seq_len: int, - ) -> torch.Tensor: - """Run action branch with cached video K/V instead of recomputing video tokens. - - Args: - action_tokens: Action tokens before layer 0, shape [B, Sa, D]. - action_freqs: Action RoPE frequencies, shape [Sa, 1, rope_dim]. - action_t_mod: Action time modulation tensor. - action_context_payload: Optional dict for action cross-attention. - - `context`: encoder states [B, L, D] - - `mask`: attention mask [B, Sa, L] or [B, 1, Sa, L] - video_kv_cache: Layer-wise cached video K/V from `prefill_video_cache`. - attention_mask: Joint [video+action] mask, shape [Sv+Sa, Sv+Sa]. - video_seq_len: Video token count `Sv` in the joint sequence prefix. - - Returns: - Updated action tokens after all layers, shape [B, Sa, D]. - """ - if "action" not in self.mixtures: - raise ValueError("MoT requires `action` expert for `forward_action_with_video_cache`.") - if len(video_kv_cache) != self.num_layers: - raise ValueError(f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}.") - if attention_mask.ndim != 2: - raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") - if attention_mask.shape[0] != attention_mask.shape[1]: - raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") - - action_seq_len = int(action_tokens.shape[1]) - total_seq_len = int(video_seq_len) + action_seq_len - if attention_mask.shape[0] != total_seq_len: - raise ValueError(f"`attention_mask` seq length mismatch: mask={attention_mask.shape[0]} vs expected_total={total_seq_len}") - # Use the action query rows from the joint [video+action] mask. - action_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] - - expert = self.mixtures["action"] - x = action_tokens - for layer_idx in range(self.num_layers): - block = expert.blocks[layer_idx] - # Action query/key/value are still step-dependent and must be recomputed each step. - ( - q_action, - k_action, - v_action, - residual_x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - use_gradient_checkpointing, - ) = self._build_expert_attention_io( - expert=expert, - block=block, - x=x, - freqs=action_freqs, - t_mod=action_t_mod, - ) - layer_cache = video_kv_cache[layer_idx] - if "k" not in layer_cache or "v" not in layer_cache: - raise ValueError(f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`.") - - k_video = layer_cache["k"] - v_video = layer_cache["v"] - if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: - raise ValueError(f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}.") - - # Mixed attention: action queries attend to cached video K/V plus current action K/V. - k_cat = torch.cat([k_video, k_action], dim=1) - v_cat = torch.cat([v_video, v_action], dim=1) - mixed = self._mixed_attention( - q_cat=q_action, - k_cat=k_cat, - v_cat=v_cat, - attention_mask=action_attention_mask, - ) - x = self._apply_post_with_optional_checkpoint( - block=block, - residual_x=residual_x, - gate_msa=gate_msa, - shift_mlp=shift_mlp, - scale_mlp=scale_mlp, - gate_mlp=gate_mlp, - use_gradient_checkpointing=use_gradient_checkpointing, - mixed_slice=mixed, - context_payload=action_context_payload, - ) - return x - - def forward( - self, - embeds_all: Dict[str, torch.Tensor], - attention_mask: torch.Tensor, - freqs_all: Dict[str, torch.Tensor], - context_all: Dict[str, Optional[dict]], - t_mod_all: Dict[str, torch.Tensor], - ): - missing = [k for k in self.expert_order if k not in embeds_all] - if missing: - raise ValueError(f"Missing expert tokens for {missing}") - missing = [k for k in self.expert_order if k not in freqs_all] - if missing: - raise ValueError(f"Missing expert freqs for {missing}") - missing = [k for k in self.expert_order if k not in t_mod_all] - if missing: - raise ValueError(f"Missing expert t_mod for {missing}") - - if attention_mask.ndim != 2: - raise ValueError(f"`attention_mask` must be 2D [S, S], got shape {tuple(attention_mask.shape)}") - if attention_mask.shape[0] != attention_mask.shape[1]: - raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") - - tokens_all = {k: v for k, v in embeds_all.items()} - - for layer_idx in range(self.num_layers): - q_chunks = [] - k_chunks = [] - v_chunks = [] - cached = {} - seq_lens = [] - - for name in self.expert_order: - expert = self.mixtures[name] - block = expert.blocks[layer_idx] - x = tokens_all[name] - freqs = freqs_all[name] - t_mod = t_mod_all[name] - - ( - q, - k, - v, - residual_x, - gate_msa, - shift_mlp, - scale_mlp, - gate_mlp, - use_gradient_checkpointing, - ) = self._build_expert_attention_io( - expert=expert, - block=block, - x=x, - freqs=freqs, - t_mod=t_mod, - ) - - q_chunks.append(q) - k_chunks.append(k) - v_chunks.append(v) - seq_lens.append(x.shape[1]) - cached[name] = { - "block": block, - "residual_x": residual_x, - "gate_msa": gate_msa, - "shift_mlp": shift_mlp, - "scale_mlp": scale_mlp, - "gate_mlp": gate_mlp, - "use_gradient_checkpointing": use_gradient_checkpointing, - } - - # 3. concat all tokens for mixed attention - q_cat = torch.cat(q_chunks, dim=1) - k_cat = torch.cat(k_chunks, dim=1) - v_cat = torch.cat(v_chunks, dim=1) - - total_seq = q_cat.shape[1] - if attention_mask.shape[0] != total_seq: - raise ValueError(f"Attention mask seq length mismatch: mask={attention_mask.shape[0]} vs tokens={total_seq}") - - mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) - - start = 0 - for name, seq_len in zip(self.expert_order, seq_lens): - # 4. split mixed attention output and apply post-attention blocks for each expert - end = start + seq_len - mixed_slice = mixed[:, start:end, :] - cached_expert = cached[name] - block = cached_expert["block"] - context_payload = context_all.get(name) - - updated_tokens = self._apply_post_with_optional_checkpoint( - block=block, - residual_x=cached_expert["residual_x"], - gate_msa=cached_expert["gate_msa"], - shift_mlp=cached_expert["shift_mlp"], - scale_mlp=cached_expert["scale_mlp"], - gate_mlp=cached_expert["gate_mlp"], - use_gradient_checkpointing=cached_expert["use_gradient_checkpointing"], - mixed_slice=mixed_slice, - context_payload=context_payload, - ) - - tokens_all[name] = updated_tokens - start = end - - return tokens_all diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py deleted file mode 100644 index ac756f593..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py +++ /dev/null @@ -1,221 +0,0 @@ -from typing import Any - -import torch -from einops import rearrange - -from ..modules.model2_2 import WanModel -from .layers import expert_block_forward, timestep_embedding - - -class FastWAMVideoDiT(WanModel): - """Wan2.2 video expert with FastWAM's split pre/post interface.""" - - def __init__( - self, - hidden_dim: int, - in_dim: int, - ffn_dim: int, - out_dim: int, - text_dim: int, - freq_dim: int, - eps: float, - patch_size: tuple[int, int, int], - num_heads: int, - attn_head_dim: int, - num_layers: int, - has_image_input: bool, - has_image_pos_emb: bool = False, - has_ref_conv: bool = False, - add_control_adapter: bool = False, - in_dim_control_adapter: int = 24, - seperated_timestep: bool = False, - require_vae_embedding: bool = False, - require_clip_embedding: bool = False, - fuse_vae_embedding_in_latents: bool = True, - action_conditioned: bool = False, - action_dim: int = 7, - action_group_causal_mask_mode: str = "causal", - video_attention_mask_mode: str = "bidirectional", - use_gradient_checkpointing: bool = False, - ): - del has_image_pos_emb, in_dim_control_adapter, action_dim, action_group_causal_mask_mode - if has_image_input or has_ref_conv or add_control_adapter: - raise ValueError("FastWAM training only supports the Wan2.2 latent-input video expert.") - if require_vae_embedding or require_clip_embedding or not fuse_vae_embedding_in_latents: - raise ValueError("FastWAM requires the first-frame VAE latent to be fused into the video input.") - if action_conditioned: - raise ValueError("FastWAM conditions actions through the MoT action expert, not the video expert.") - if hidden_dim != num_heads * attn_head_dim: - raise ValueError(f"hidden_dim must equal num_heads * attn_head_dim, got {hidden_dim} and {num_heads} * {attn_head_dim}") - - super().__init__( - model_type="ti2v", - patch_size=tuple(patch_size), - text_len=512, - in_dim=in_dim, - dim=hidden_dim, - ffn_dim=ffn_dim, - freq_dim=freq_dim, - text_dim=text_dim, - out_dim=out_dim, - num_heads=num_heads, - num_layers=num_layers, - window_size=(-1, -1), - qk_norm=True, - cross_attn_norm=True, - eps=eps, - ) - self.hidden_dim = hidden_dim - self.attn_head_dim = attn_head_dim - self.seperated_timestep = seperated_timestep - self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents - self.video_attention_mask_mode = str(video_attention_mask_mode) - self.use_gradient_checkpointing = bool(use_gradient_checkpointing) - - def _validate_inputs(self, x, timestep, context, context_mask): - if x.ndim != 5: - raise ValueError(f"latents must be [B,C,T,H,W], got {tuple(x.shape)}") - if context.ndim != 3: - raise ValueError(f"context must be [B,L,D], got {tuple(context.shape)}") - if timestep.ndim != 1: - raise ValueError(f"timestep must be [B] or [1], got {tuple(timestep.shape)}") - batch_size = x.shape[0] - if self.training and batch_size != 1: - raise ValueError("FastWAM video training only supports physical batch size 1.") - if context.shape[0] != batch_size: - if not self.training and batch_size == 1: - x = x.expand(context.shape[0], -1, -1, -1, -1) - batch_size = context.shape[0] - else: - raise ValueError(f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}") - if timestep.shape[0] not in (1, batch_size): - raise ValueError(f"timestep length must be 1 or {batch_size}, got {timestep.shape[0]}") - if timestep.shape[0] == 1 and batch_size > 1: - timestep = timestep.expand(batch_size) - if context_mask is None: - context_mask = torch.ones(context.shape[:2], dtype=torch.bool, device=context.device) - elif context_mask.shape != context.shape[:2]: - raise ValueError(f"context_mask must have shape {tuple(context.shape[:2])}, got {tuple(context_mask.shape)}") - return x, timestep, context_mask - - def build_video_to_video_mask(self, video_seq_len, video_tokens_per_frame, device): - if self.video_attention_mask_mode == "bidirectional": - return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) - if video_seq_len % video_tokens_per_frame != 0: - raise ValueError("Video token count must be divisible by tokens per frame.") - if self.video_attention_mask_mode == "per_frame_causal": - frame_count = video_seq_len // video_tokens_per_frame - mask = torch.tril(torch.ones((frame_count, frame_count), dtype=torch.bool, device=device)) - return mask.repeat_interleave(video_tokens_per_frame, 0).repeat_interleave(video_tokens_per_frame, 1) - if self.video_attention_mask_mode == "first_frame_causal": - mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) - mask[:video_tokens_per_frame, video_tokens_per_frame:] = False - return mask - raise ValueError(f"Unsupported video_attention_mask_mode: {self.video_attention_mask_mode}") - - def _video_frequencies(self, grid_size, device): - frames, height, width = grid_size - half_dim = self.attn_head_dim // 2 - split_sizes = [half_dim - 2 * (half_dim // 3), half_dim // 3, half_dim // 3] - time_freqs, height_freqs, width_freqs = self.freqs.split(split_sizes, dim=1) - frequencies = torch.cat( - [ - time_freqs[:frames].view(frames, 1, 1, -1).expand(frames, height, width, -1), - height_freqs[:height].view(1, height, 1, -1).expand(frames, height, width, -1), - width_freqs[:width].view(1, 1, width, -1).expand(frames, height, width, -1), - ], - dim=-1, - ) - return frequencies.reshape(frames * height * width, 1, -1).to(device) - - def pre_dit( - self, - x: torch.Tensor, - timestep: torch.Tensor, - context: torch.Tensor, - context_mask: torch.Tensor | None = None, - action: torch.Tensor | None = None, - fuse_vae_embedding_in_latents: bool = False, - control_camera_latents_input: torch.Tensor | None = None, - ) -> dict[str, Any]: - del action, control_camera_latents_input - x, timestep, context_mask = self._validate_inputs(x, timestep, context, context_mask) - if not self.seperated_timestep or not fuse_vae_embedding_in_latents: - raise ValueError("FastWAM video expert requires separated timestep with fused first-frame latents.") - - batch_size = x.shape[0] - patch_h, patch_w = self.patch_size[1:] - if x.shape[3] % patch_h or x.shape[4] % patch_w: - raise ValueError("Latent spatial dimensions must be divisible by the Wan patch size.") - tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) - token_timesteps = timestep.view(batch_size, 1, 1).expand(batch_size, x.shape[2], tokens_per_frame).clone() - token_timesteps[:, 0] = 0 - token_timesteps = token_timesteps.reshape(batch_size, -1) - time_embedding = self.time_embedding(timestep_embedding(self.freq_dim, token_timesteps.reshape(-1))).reshape(batch_size, -1, self.hidden_dim) - time_modulation = self.time_projection(time_embedding).unflatten(2, (6, self.hidden_dim)) - - patches = self.patch_embedding(x) - frames, height, width = patches.shape[2:] - tokens = rearrange(patches, "b c f h w -> b (f h w) c").contiguous() - context = self.text_embedding(context) - context_mask = context_mask.unsqueeze(1).expand(-1, tokens.shape[1], -1) - return { - "tokens": tokens, - "freqs": self._video_frequencies((frames, height, width), tokens.device), - "t": time_embedding, - "t_mod": time_modulation, - "context": context, - "context_mask": context_mask, - "meta": { - "grid_size": (frames, height, width), - "tokens_per_frame": tokens_per_frame, - "batch_size": batch_size, - }, - } - - def post_dit(self, tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: - time_embedding = pre_state["t"] - shift, scale = (self.head.modulation.unsqueeze(0).to(time_embedding) + time_embedding.unsqueeze(2)).chunk(2, dim=2) - tokens = self.head.head(self.head.norm(tokens) * (1 + scale.squeeze(2)) + shift.squeeze(2)) - frames, height, width = pre_state["meta"]["grid_size"] - return rearrange( - tokens, - "b (f h w) (pf ph pw c) -> b c (f pf) (h ph) (w pw)", - f=frames, - h=height, - w=width, - pf=self.patch_size[0], - ph=self.patch_size[1], - pw=self.patch_size[2], - c=self.out_dim, - ) - - def forward(self, x, timestep, context, context_mask=None, action=None, fuse_vae_embedding_in_latents=False): - state = self.pre_dit( - x=x, - timestep=timestep, - context=context, - context_mask=context_mask, - action=action, - fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, - ) - tokens = state["tokens"] - self_mask = self.build_video_to_video_mask(tokens.shape[1], state["meta"]["tokens_per_frame"], tokens.device) - for block in self.blocks: - - def run(value, current_block=block): - return expert_block_forward( - current_block, - value, - state["context"], - state["t_mod"], - state["freqs"], - state["context_mask"], - self_mask, - ) - - if self.use_gradient_checkpointing and self.training: - tokens = torch.utils.checkpoint.checkpoint(run, tokens, use_reentrant=False) - else: - tokens = run(tokens) - return self.post_dit(tokens, state) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py b/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py deleted file mode 100644 index b1d527b9e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/capability_adapters/wan_world_action_capability.py +++ /dev/null @@ -1,40 +0,0 @@ -"""World-action training capability for Wan FastWAM.""" - -from lightx2v_train.model_capabilities import ( - BoundCapability, - LossResult, - WorldActionTrainingCapability, -) - - -class WanWorldActionCapability( - BoundCapability, - WorldActionTrainingCapability, -): - def configure(self) -> None: - self.model.set_dit_only_trainable() - self.model.log_model_structure() - - def parameters(self): - return self.model.trainable_parameters() - - def module(self): - return self.model.unwrap_module() - - def compute_loss(self, batch, module=None): - if module is None: - module = self.module() - with self.model.autocast_context(): - loss, metrics = module(batch) - return LossResult(loss=loss, metrics=metrics) - - def evaluation_loss(self, batch): - with self.model.autocast_context(): - loss, metrics = self.module().training_loss(batch) - return LossResult(loss=loss, metrics=metrics) - - def load_checkpoint(self, path) -> None: - self.model.load_checkpoint(path) - - def save_checkpoint(self, path, step=None) -> None: - self.model.save_checkpoint(path, step=step) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py deleted file mode 100644 index 67a7af58a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/wan/wan_fastwam.py +++ /dev/null @@ -1,129 +0,0 @@ -import os -from contextlib import nullcontext - -import torch -from loguru import logger - -from lightx2v_train.model_capabilities import ( - CapabilityProvider, - WorldActionTrainingCapability, -) -from lightx2v_train.model_zoo.native.wan.fastwam import FastWAM -from lightx2v_train.model_zoo.wan.capability_adapters.wan_world_action_capability import ( - WanWorldActionCapability, -) -from lightx2v_train.utils.registry import MODEL_REGISTER -from lightx2v_train.utils.utils import get_running_dtype - - -def _resolve_local_path(path, name, directory=False): - if not path: - raise ValueError(f"model.{name} must be an explicit local path.") - resolved = os.path.abspath(os.path.expanduser(str(path))) - exists = os.path.isdir(resolved) if directory else os.path.isfile(resolved) - if not exists: - expected = "directory" if directory else "file" - raise FileNotFoundError(f"model.{name} {expected} does not exist: {resolved}") - return resolved - - -@MODEL_REGISTER("wan_fastwam") -class WanFastWAMModel(CapabilityProvider): - def __init__(self, config): - super().__init__() - self.config = config - self.model_config = config["model"] - self.running_dtype = get_running_dtype(self.model_config.get("running_dtype", "bf16")) - self.device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") - self.module = None - - def register_capabilities(self): - self.capabilities.register( - WorldActionTrainingCapability, - WanWorldActionCapability(self), - ) - - def load_components( - self, - *, - load_transformer, - load_vae, - load_condition_encoder, - ): - if not (load_transformer and load_vae and load_condition_encoder): - raise ValueError("wan_fastwam does not support selective component loading.") - - model_path = _resolve_local_path(self.model_config.get("model_path"), "model_path", directory=True) - configured_action_dit_path = self.model_config.get("action_dit_pretrained_path") - action_dit_pretrained_path = _resolve_local_path(configured_action_dit_path, "action_dit_pretrained_path") if configured_action_dit_path else None - load_text_encoder = bool(self.model_config.get("load_text_encoder", False)) - - video_scheduler = self.model_config.get("video_scheduler", {}) - action_scheduler = self.model_config.get("action_scheduler", {}) - loss_config = self.model_config.get("loss", {}) - self.module = FastWAM.from_wan22_pretrained( - model_path=model_path, - device=str(self.device), - torch_dtype=self.running_dtype, - tokenizer_max_len=int(self.model_config.get("tokenizer_max_len", 128)), - load_text_encoder=load_text_encoder, - proprio_dim=self.model_config.get("proprio_dim", 8), - video_dit_config=dict(self.model_config.get("video_dit_config", {})), - action_dit_config=dict(self.model_config.get("action_dit_config", {})), - action_dit_pretrained_path=action_dit_pretrained_path, - skip_dit_load_from_pretrain=bool(self.model_config.get("skip_dit_load_from_pretrain", False)), - mot_checkpoint_mixed_attn=bool(self.model_config.get("mot_checkpoint_mixed_attn", False)), - video_train_shift=float(video_scheduler.get("train_shift", 5.0)), - video_infer_shift=float(video_scheduler.get("infer_shift", 5.0)), - video_num_train_timesteps=int(video_scheduler.get("num_train_timesteps", 1000)), - action_train_shift=float(action_scheduler.get("train_shift", 5.0)), - action_infer_shift=float(action_scheduler.get("infer_shift", 5.0)), - action_num_train_timesteps=int(action_scheduler.get("num_train_timesteps", 1000)), - loss_lambda_video=float(loss_config.get("lambda_video", 1.0)), - loss_lambda_action=float(loss_config.get("lambda_action", 1.0)), - ) - - checkpoint_path = self.model_config.get("checkpoint_path") - if checkpoint_path: - logger.info("[model] loading FastWAM checkpoint {}", checkpoint_path) - self.unwrap_module().load_checkpoint(checkpoint_path) - - def unwrap_module(self): - if self.module is None: - raise RuntimeError("FastWAM module has not been loaded.") - return self.module - - def autocast_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 set_dit_only_trainable(self): - module = self.unwrap_module() - module.eval() - module.requires_grad_(False) - module.dit.train() - module.dit.requires_grad_(True) - proprio_encoder = getattr(module, "proprio_encoder", None) - if proprio_encoder is not None: - proprio_encoder.train() - proprio_encoder.requires_grad_(True) - - def trainable_parameters(self): - module = self.unwrap_module() - params = list(module.dit.parameters()) - proprio_encoder = getattr(module, "proprio_encoder", None) - if proprio_encoder is not None: - params.extend(list(proprio_encoder.parameters())) - return (param for param in params if param.requires_grad) - - def save_checkpoint(self, path, step=None): - self.unwrap_module().save_checkpoint(path, step=step) - - def load_checkpoint(self, path): - return self.unwrap_module().load_checkpoint(path) - - def log_model_structure(self): - module = self.unwrap_module() - logger.info("[model] class=WanFastWAMModel fastwam_class={}", module.__class__.__name__) - logger.info("[model] trainable root={}", module.dit.__class__.__name__) diff --git a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py index e3ee864ca..08d51eb99 100644 --- a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py +++ b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py @@ -219,67 +219,3 @@ def sample_clean_augmentation(self, num_frames, num_frame_per_chunk, max_timeste def add_noise(self, latent, noise, sigmas): sigmas = sigmas.reshape(sigmas.shape[0], 1, sigmas.shape[1], 1, 1) return (1.0 - sigmas) * latent + sigmas * noise - - -class WanContinuousFlowMatchScheduler: - """Continuous flow matching used by FastWAM video and action experts.""" - - def __init__(self, num_train_timesteps=1000, shift=5.0, eps=1e-10): - if num_train_timesteps <= 0: - raise ValueError("num_train_timesteps must be positive.") - if shift <= 0: - raise ValueError("shift must be positive.") - self.num_train_timesteps = int(num_train_timesteps) - self.shift = float(shift) - self.eps = float(eps) - self._y_min, self._weight_norm_const = self._precompute_training_weight_stats() - - @staticmethod - def _phi(value, shift): - return shift * value / (1.0 + (shift - 1.0) * value) - - def _precompute_training_weight_stats(self): - steps = self.num_train_timesteps - grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1] - timestep = self._phi(grid, self.shift) * float(steps) - weights = torch.exp(-2.0 * ((timestep - steps / 2.0) / steps) ** 2) - minimum = float(weights.min().item()) - return minimum, float((weights - minimum).mean().item()) - - def sample_training_t(self, device, dtype): - value = torch.rand((1,), device=device, dtype=torch.float32) - return (self._phi(value, self.shift) * self.num_train_timesteps).to(dtype=dtype) - - def training_weight(self, timestep): - timestep = timestep.to(dtype=torch.float32) - steps = float(self.num_train_timesteps) - weight = torch.exp(-2.0 * ((timestep - steps / 2.0) / steps) ** 2) - weight = (weight - self._y_min) / (self._weight_norm_const + self.eps) - return weight.reshape(()) if weight.numel() == 1 else weight - - def add_noise(self, original_samples, noise, timestep): - sigma = (timestep / self.num_train_timesteps).to(device=original_samples.device, dtype=original_samples.dtype) - if sigma.ndim: - sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) - return (1 - sigma) * original_samples + sigma * noise - - @staticmethod - def training_target(sample, noise, timestep): - del timestep - return noise - sample - - def build_inference_schedule(self, num_inference_steps, device, dtype, shift_override=None): - if num_inference_steps <= 0: - raise ValueError("num_inference_steps must be positive.") - shift = self.shift if shift_override is None else float(shift_override) - values = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device) - sigmas = self._phi(values, shift) - timesteps = sigmas[:-1] * self.num_train_timesteps - return timesteps.to(dtype=dtype), (sigmas[1:] - sigmas[:-1]).to(dtype=dtype) - - @staticmethod - def step(model_output, delta, sample): - delta = delta.to(device=sample.device, dtype=sample.dtype) - if delta.ndim: - delta = delta.view(-1, *([1] * (sample.ndim - 1))) - return sample + model_output * delta diff --git a/lightx2v_train/lightx2v_train/trainers/__init__.py b/lightx2v_train/lightx2v_train/trainers/__init__.py index 75aa1d806..1524f4620 100644 --- a/lightx2v_train/lightx2v_train/trainers/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/__init__.py @@ -30,7 +30,6 @@ ".training_cache", "TrainingCacheTrainer", ), - "WorldActionTrainer": (".world_action", "WorldActionTrainer"), } diff --git a/lightx2v_train/lightx2v_train/trainers/world_action.py b/lightx2v_train/lightx2v_train/trainers/world_action.py deleted file mode 100644 index 819e8ca01..000000000 --- a/lightx2v_train/lightx2v_train/trainers/world_action.py +++ /dev/null @@ -1,689 +0,0 @@ -import json -import os -import shutil -import time -from contextlib import nullcontext - -import numpy as np -import torch -from PIL import Image, ImageDraw -from loguru import logger -from torch.nn.parallel import DistributedDataParallel -from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, SequentialLR - -from lightx2v_train.data.utils import require_singleton_dataloader -from lightx2v_train.model_capabilities import WorldActionTrainingCapability -from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints -from lightx2v_train.runtime.distributed import ( - barrier, - get_data_parallel_group, - get_data_parallel_world_size, - get_rank, - get_world_size, - is_distributed, - is_main_process, - reduce_mean, -) -from lightx2v_train.utils.registry import TRAINER_REGISTER -from lightx2v_train.utils.video import pil_frames_to_video_tensor, save_mp4, video_psnr, video_ssim - -ZERO1_PER_RANK_FORMAT = "zero1_per_rank_v1" -FULL_OPTIMIZER_FORMAT = "full_v1" - - -def _optimizer_shard_filename(rank, world_size): - return f"optimizer-rank-{rank:05d}-of-{world_size:05d}.pt" - - -def _temporary_path(path): - return f"{path}.tmp-{os.getpid()}-{get_rank()}" - - -def _atomic_torch_save(payload, path): - temporary_path = _temporary_path(path) - try: - torch.save(payload, temporary_path) - os.replace(temporary_path, path) - finally: - if os.path.exists(temporary_path): - os.remove(temporary_path) - - -def _sync_optimizer_group_options(source_groups, target_groups): - if len(source_groups) != len(target_groups): - raise RuntimeError(f"Optimizer parameter-group count mismatch: source={len(source_groups)} target={len(target_groups)}") - for source, target in zip(source_groups, target_groups): - for key, value in source.items(): - if key != "params": - target[key] = value - - -def _unwrap_dataset(dataset): - while hasattr(dataset, "dataset"): - dataset = dataset.dataset - return dataset - - -@TRAINER_REGISTER("world_action") -class WorldActionTrainer: - required_capabilities = (WorldActionTrainingCapability,) - - def __init__(self, config): - self.config = config - self.model_config = config["model"] - self.training_config = config["training"] - self.inference_config = config.get("inference", {}) - self.logging_config = config.get("logging", {}) - - optimizer_config = self.training_config.get("optimizer", {}) - self.learning_rate = float(optimizer_config.get("learning_rate", self.training_config.get("learning_rate", 1e-4))) - self.weight_decay = float(optimizer_config.get("weight_decay", self.training_config.get("weight_decay", 0.0))) - self.adam_beta1 = float(optimizer_config.get("adam_beta1", 0.9)) - self.adam_beta2 = float(optimizer_config.get("adam_beta2", 0.95)) - self.adam_epsilon = float(optimizer_config.get("adam_epsilon", 1e-8)) - - self.output_train_dir = self.training_config["output_dir"] - self.max_train_iters = self._resolve_max_train_iters() - self.gradient_accumulation_iters = int(self.training_config.get("gradient_accumulation_iters", 1)) - self.max_grad_norm = float(self.training_config.get("max_grad_norm", 1.0)) - self.save_every_iters = int(self.training_config.get("save_every_iters", 0) or 0) - self.save_total_limit = int(self.training_config.get("save_total_limit", 3)) - self.save_final = bool(self.training_config.get("save_final", True)) - self.lr_scheduler_name = self.training_config.get("lr_scheduler", "constant") - self.lr_warmup_iters = int(self.training_config.get("lr_warmup_iters", 0)) - self.lr_eta_min_ratio = float(self.training_config.get("lr_eta_min_ratio", 0.01)) - self.train_log_every_iters = max(1, int(self.logging_config.get("train_log_every_iters", 10))) - - zero1_config = self.config.get("distributed", {}).get("zero1", {}) - self.zero1_requested = bool(zero1_config.get("enabled", False)) if isinstance(zero1_config, dict) else bool(zero1_config) - self.zero1_enabled = False - - self.eval_every_iters = int(self.inference_config.get("infer_every_iters", 0) or 0) - self.eval_num_inference_steps = int(self.inference_config.get("eval_num_inference_steps", 5)) - self.eval_num_samples = max(1, int(self.inference_config.get("num_samples", 1))) - self.eval_sample_strategy = str(self.inference_config.get("sample_strategy", "fixed")).lower() - self.eval_run_inference = bool(self.inference_config.get("run_inference", True)) - self.eval_save_video = bool(self.inference_config.get("save_video", True)) - self.eval_save_preview = bool(self.inference_config.get("save_preview", True)) - self.eval_fps = max(1, int(self.inference_config.get("fps", 8))) - self.eval_tiled = bool(self.inference_config.get("tiled", False)) - self.eval_output_dir = self.inference_config.get("output_dir", os.path.join(self.output_train_dir, "eval")) - self.eval_seed = int(self.inference_config.get("seed", 42)) - - resume_config = self.config.get("resume", {}) - self.auto_resume = bool(resume_config.get("auto_resume", False)) - self.resume_ckpt_path = resume_config.get("resume_ckpt_path") - - def _resolve_max_train_iters(self): - max_train_iters = self.training_config.get("max_train_iters", self.training_config.get("max_steps")) - if max_train_iters is not None: - return int(max_train_iters) - num_epochs = int(self.training_config.get("num_epochs", 1)) - return max(1, num_epochs) - - def set_model(self, model): - self.model = model - registry = model.ensure_capabilities() - self.world_action = registry.require(WorldActionTrainingCapability) - - def set_data(self, dataloader_train, dataloader_eval=None): - require_singleton_dataloader(dataloader_train, "Training dataloader") - if dataloader_eval is not None: - require_singleton_dataloader(dataloader_eval, "Evaluation dataloader") - self.dataloader_train = dataloader_train - self.dataloader_eval = dataloader_eval - - def _build_optimizer(self): - optimizer_kwargs = { - "lr": self.learning_rate, - "weight_decay": self.weight_decay, - "betas": (self.adam_beta1, self.adam_beta2), - "eps": self.adam_epsilon, - } - data_parallel_size = get_data_parallel_world_size() - if self.zero1_requested and data_parallel_size > 1: - from torch.distributed.optim import ZeroRedundancyOptimizer - - self.zero1_enabled = True - logger.info("[optimizer] using ZeroRedundancyOptimizer with AdamW across {} data-parallel ranks", data_parallel_size) - return ZeroRedundancyOptimizer( - self.trainable_params, - optimizer_class=torch.optim.AdamW, - process_group=get_data_parallel_group(), - parameters_as_bucket_view=False, - overlap_with_ddp=False, - **optimizer_kwargs, - ) - - self.zero1_enabled = False - if self.zero1_requested: - logger.info("[optimizer] ZeroRedundancyOptimizer requested but data-parallel size is 1; using AdamW") - else: - logger.info("[optimizer] using AdamW without optimizer-state sharding") - return torch.optim.AdamW(self.trainable_params, **optimizer_kwargs) - - def _build_lr_scheduler(self): - # LinearLR changes the optimizer's current LR during construction. Capture - # the configured base LR first so cosine eta_min is not scaled by warmup. - base_lr = float(self.optimizer.param_groups[0]["lr"]) - remaining_iters = self.max_train_iters - self.lr_warmup_iters - if self.lr_scheduler_name == "cosine": - main_scheduler = CosineAnnealingLR( - self.optimizer, - T_max=remaining_iters, - eta_min=base_lr * self.lr_eta_min_ratio, - ) - else: - main_scheduler = ConstantLR( - self.optimizer, - factor=1.0, - total_iters=remaining_iters, - ) - - if self.lr_warmup_iters == 0: - return main_scheduler - - warmup_scheduler = LinearLR( - self.optimizer, - start_factor=1.0 / self.lr_warmup_iters, - end_factor=1.0, - total_iters=self.lr_warmup_iters, - ) - return SequentialLR( - self.optimizer, - schedulers=[warmup_scheduler, main_scheduler], - milestones=[self.lr_warmup_iters], - ) - - def setup(self, resume_ckpt_path=None): - self.world_action.configure() - self.trainable_params = list(self.world_action.parameters()) - if not self.trainable_params: - raise RuntimeError("FastWAM has no trainable parameters.") - - self.optimizer = self._build_optimizer() - self.lr_scheduler = self._build_lr_scheduler() - self.train_module = self.world_action.module() - self.ddp_enabled = False - if is_distributed(): - self.train_module = DistributedDataParallel( - self.train_module, - device_ids=[torch.cuda.current_device()] if torch.cuda.is_available() else None, - find_unused_parameters=bool(self.training_config.get("ddp_find_unused_parameters", False)), - ) - self.ddp_enabled = True - - if resume_ckpt_path is not None: - self._load_resume_state(resume_ckpt_path) - - def _resolve_resume(self): - if self.resume_ckpt_path: - current_iter = parse_checkpoint_iteration(self.resume_ckpt_path) - return self.resume_ckpt_path, current_iter - if not self.auto_resume: - return None, 0 - ckpt_path, current_iter = find_latest_checkpoint(self.output_train_dir) - if ckpt_path is None: - logger.info("Auto-resume enabled but no checkpoint found in '{}'. Starting from scratch.", self.output_train_dir) - return None, 0 - logger.info("Auto-resuming from checkpoint: {} (iteration {})", ckpt_path, current_iter) - return ckpt_path, current_iter - - def _load_resume_state(self, resume_ckpt_path): - weights_path = os.path.join(resume_ckpt_path, "fastwam.pt") - state_path = os.path.join(resume_ckpt_path, "training_state.pt") - if not os.path.exists(weights_path): - raise RuntimeError(f"fastwam.pt not found in {resume_ckpt_path}") - if not os.path.exists(state_path): - raise RuntimeError(f"training_state.pt not found in {resume_ckpt_path}") - self.world_action.load_checkpoint(weights_path) - state = torch.load(state_path, map_location="cpu", weights_only=False) - if state.get("world_size") != get_world_size(): - raise RuntimeError(f"Cannot resume world_size={state.get('world_size')} checkpoint with world_size={get_world_size()}.") - optimizer_state_format = state.get("optimizer_state_format") - expected_optimizer_sharding = "zero1" if self.zero1_enabled else "none" - if state.get("optimizer_sharding") != expected_optimizer_sharding: - raise RuntimeError(f"Optimizer sharding mode mismatch: checkpoint={state.get('optimizer_sharding')!r} active={expected_optimizer_sharding!r}.") - if optimizer_state_format == ZERO1_PER_RANK_FORMAT: - if not self.zero1_enabled: - raise RuntimeError("A per-rank ZeRO-1 optimizer checkpoint requires distributed.zero1.enabled=true.") - shard_files = state.get("optimizer_shard_files") - if not isinstance(shard_files, list) or len(shard_files) != get_world_size(): - raise RuntimeError(f"Invalid ZeRO-1 optimizer shard manifest: expected {get_world_size()} files, got {shard_files!r}.") - shard_path = os.path.join(resume_ckpt_path, shard_files[get_rank()]) - if not os.path.isfile(shard_path): - raise RuntimeError(f"ZeRO-1 optimizer shard not found for rank {get_rank()}: {shard_path}") - shard = torch.load(shard_path, map_location="cpu", weights_only=False) - if shard.get("rank") != get_rank() or shard.get("world_size") != get_world_size(): - raise RuntimeError( - f"ZeRO-1 optimizer shard metadata mismatch: rank={shard.get('rank')} world_size={shard.get('world_size')} expected_rank={get_rank()} expected_world_size={get_world_size()}." - ) - self.optimizer.optim.load_state_dict(shard["optimizer"]) - _sync_optimizer_group_options(self.optimizer.optim.param_groups, self.optimizer.param_groups) - elif optimizer_state_format == FULL_OPTIMIZER_FORMAT: - if self.zero1_enabled: - raise RuntimeError("A full optimizer checkpoint requires distributed.zero1.enabled=false.") - self.optimizer.load_state_dict(state["optimizer"]) - else: - raise RuntimeError(f"Unsupported optimizer checkpoint format {optimizer_state_format!r} in {state_path}.") - self.lr_scheduler.load_state_dict(state["lr_scheduler"]) - logger.info( - "[resume] restored FastWAM training state from {} (optimizer sharding: {}, state format: {})", - resume_ckpt_path, - expected_optimizer_sharding, - optimizer_state_format, - ) - - def _save_checkpoint(self, iteration): - save_start_time = time.perf_counter() - if is_main_process(): - prune_checkpoints(self.output_train_dir, self.save_total_limit) - save_dir = os.path.join(self.output_train_dir, f"checkpoint-{iteration:09d}") - logger.info("[checkpoint] saving FastWAM iter={} path={}", iteration, save_dir) - if is_main_process(): - os.makedirs(save_dir, exist_ok=True) - barrier() - - optimizer_state_format = FULL_OPTIMIZER_FORMAT - optimizer_shard_files = None - if self.zero1_enabled: - optimizer_state_format = ZERO1_PER_RANK_FORMAT - optimizer_shard_files = [_optimizer_shard_filename(rank, get_world_size()) for rank in range(get_world_size())] - _sync_optimizer_group_options(self.optimizer.param_groups, self.optimizer.optim.param_groups) - shard_path = os.path.join(save_dir, optimizer_shard_files[get_rank()]) - shard_start_time = time.perf_counter() - _atomic_torch_save( - { - "rank": get_rank(), - "world_size": get_world_size(), - "optimizer": self.optimizer.optim.state_dict(), - }, - shard_path, - ) - logger.info( - "[checkpoint] saved optimizer shard rank={} duration={:.3f}s path={}", - get_rank(), - time.perf_counter() - shard_start_time, - shard_path, - ) - barrier() - - if is_main_process(): - weights_path = os.path.join(save_dir, "fastwam.pt") - temporary_weights_path = _temporary_path(weights_path) - try: - self.world_action.save_checkpoint( - temporary_weights_path, - step=iteration, - ) - os.replace(temporary_weights_path, weights_path) - finally: - if os.path.exists(temporary_weights_path): - os.remove(temporary_weights_path) - - training_state = { - "iteration": iteration, - "world_size": get_world_size(), - "optimizer_sharding": "zero1" if self.zero1_enabled else "none", - "optimizer_state_format": optimizer_state_format, - "lr_scheduler": self.lr_scheduler.state_dict(), - } - if self.zero1_enabled: - training_state["optimizer_shard_files"] = optimizer_shard_files - else: - training_state["optimizer"] = self.optimizer.state_dict() - _atomic_torch_save(training_state, os.path.join(save_dir, "training_state.pt")) - config_path = self.config.get("config_path") - if config_path is not None: - shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) - barrier() - logger.info( - "[checkpoint] saved FastWAM iter={} duration={:.3f}s path={}", - iteration, - time.perf_counter() - save_start_time, - save_dir, - ) - - def _forward_loss(self, sample): - result = self.world_action.compute_loss( - sample, - module=self.train_module, - ) - return result.loss, result.metrics - - def _ddp_sync_context(self, sync_grad): - if self.ddp_enabled and not sync_grad: - return self.train_module.no_sync() - return nullcontext() - - def train(self): - resume_ckpt_path, current_iter = self._resolve_resume() - self.setup(resume_ckpt_path=resume_ckpt_path) - if is_main_process(): - os.makedirs(self.output_train_dir, exist_ok=True) - barrier() - - grad_accum_counter = 0 - running_loss = 0.0 - running_loss_dict = {} - epoch = 0 - run_start_iter = current_iter - start_time = time.perf_counter() - - logger.info( - "[train] start method=world_action iter={}/{} world_size={} grad_accum={} log_every={}", - current_iter, - self.max_train_iters, - get_world_size(), - self.gradient_accumulation_iters, - self.train_log_every_iters, - ) - - while current_iter < self.max_train_iters: - sampler = getattr(self.dataloader_train, "sampler", None) - if hasattr(sampler, "set_epoch"): - sampler.set_epoch(epoch) - - for sample in self.dataloader_train: - sync_grad = (grad_accum_counter + 1) % self.gradient_accumulation_iters == 0 - with self._ddp_sync_context(sync_grad): - loss, loss_dict = self._forward_loss(sample) - (loss / self.gradient_accumulation_iters).backward() - - running_loss += float(loss.detach().item()) / self.gradient_accumulation_iters - for key, value in loss_dict.items(): - running_loss_dict[key] = running_loss_dict.get(key, 0.0) + float(value) / self.gradient_accumulation_iters - - grad_accum_counter += 1 - if grad_accum_counter % self.gradient_accumulation_iters != 0: - continue - - grad_norm = torch.nn.utils.clip_grad_norm_(self.trainable_params, self.max_grad_norm) - self.optimizer.step() - self.lr_scheduler.step() - self.optimizer.zero_grad(set_to_none=True) - - current_iter += 1 - display_loss = reduce_mean(running_loss) - current_lr = self.lr_scheduler.get_last_lr()[0] - if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= self.max_train_iters: - elapsed = max(time.perf_counter() - start_time, 1e-6) - detail = " ".join(f"{key}={reduce_mean(value):.6f}" for key, value in sorted(running_loss_dict.items())) - logger.info( - "[train] iter={}/{} loss={:.6f} {} lr={:.8f} grad_norm={:.4f} speed={:.3f} it/s", - current_iter, - self.max_train_iters, - display_loss, - detail, - current_lr, - float(grad_norm), - (current_iter - run_start_iter) / elapsed, - ) - running_loss = 0.0 - running_loss_dict = {} - - if self.eval_every_iters > 0 and self.dataloader_eval is not None and current_iter % self.eval_every_iters == 0: - self.evaluate(current_iter) - - if self.save_every_iters > 0 and current_iter % self.save_every_iters == 0: - self._save_checkpoint(current_iter) - - if current_iter >= self.max_train_iters: - break - - epoch += 1 - - if self.save_final and (self.save_every_iters <= 0 or current_iter % self.save_every_iters != 0): - self._save_checkpoint(current_iter) - logger.info("[train] finished FastWAM iter={}/{}", current_iter, self.max_train_iters) - - @torch.no_grad() - def evaluate(self, current_iter): - if self.dataloader_eval is None: - return - - if torch.cuda.is_available(): - torch.cuda.synchronize() - eval_start_time = time.perf_counter() - eval_dataset = self.dataloader_eval.dataset - model_dataset = _unwrap_dataset(eval_dataset) - if len(eval_dataset) == 0: - return - - model = self.world_action.module() - model.eval() - step_dir = os.path.join(self.eval_output_dir, f"step_{current_iter:09d}") - os.makedirs(step_dir, exist_ok=True) - - sample_selection_seed = self._eval_sample_selection_seed(current_iter) - eval_indices = self._select_eval_indices(len(eval_dataset), sample_selection_seed) - local_results = [] - for sample_number in range(get_rank(), len(eval_indices), get_world_size()): - eval_index = eval_indices[sample_number] - sample = self._to_batched_eval_sample(eval_dataset[eval_index]) - local_results.append( - self._evaluate_sample( - model=model, - dataset=model_dataset, - sample=sample, - sample_number=sample_number, - eval_index=eval_index, - step_dir=step_dir, - ) - ) - - results = self._gather_eval_results(local_results) - summary = self._summarize_eval_results(results) - if torch.cuda.is_available(): - torch.cuda.synchronize() - duration_seconds = time.perf_counter() - eval_start_time - if is_main_process(): - metrics_path = os.path.join(step_dir, "metrics.json") - with open(metrics_path, "w", encoding="utf-8") as handle: - json.dump( - { - "iteration": int(current_iter), - "num_samples": len(results), - "sample_strategy": self.eval_sample_strategy, - "sample_selection_seed": sample_selection_seed, - "duration_seconds": duration_seconds, - "mean": summary, - "samples": results, - }, - handle, - ensure_ascii=True, - indent=2, - ) - details = " ".join(f"{key}={value:.6f}" for key, value in sorted(summary.items())) - logger.info( - "[eval] iter={} samples={} duration={:.3f}s {} output={}", - current_iter, - len(results), - duration_seconds, - details, - step_dir, - ) - - barrier() - self.world_action.configure() - - def _to_batched_eval_sample(self, sample): - result = {} - for key, value in sample.items(): - if isinstance(value, torch.Tensor): - result[key] = value.unsqueeze(0) if value.ndim >= 1 else value.reshape(1) - elif isinstance(value, str): - result[key] = [value] - else: - result[key] = value - return result - - def _eval_sample_selection_seed(self, current_iter): - if self.eval_sample_strategy == "rotating": - eval_round = current_iter // max(self.eval_every_iters, 1) - return self.eval_seed + eval_round - return self.eval_seed - - def _select_eval_indices(self, dataset_size, sample_selection_seed): - count = min(self.eval_num_samples, dataset_size) - generator = torch.Generator(device="cpu").manual_seed(sample_selection_seed) - return torch.randperm(dataset_size, generator=generator)[:count].tolist() - - def _evaluate_sample(self, model, dataset, sample, sample_number, eval_index, step_dir): - sample_seed = self.eval_seed + sample_number - rng_devices = [torch.cuda.current_device()] if torch.cuda.is_available() else [] - with torch.random.fork_rng(devices=rng_devices): - torch.manual_seed(sample_seed) - evaluation = self.world_action.evaluation_loss(sample) - val_loss = evaluation.loss - val_loss_dict = evaluation.metrics - result = { - "sample_number": int(sample_number), - "dataset_index": int(eval_index), - "val_loss": float(val_loss.detach().float().item()), - } - result.update({f"val_{key}": float(value) for key, value in val_loss_dict.items()}) - if not self.eval_run_inference: - return result - - video0 = sample["video"][0] - action = sample.get("action") - proprio = sample.get("proprio") - context = sample.get("context") - context_mask = sample.get("context_mask") - prompt = sample.get("prompt", [None])[0] - input_image = video0[:, 0].unsqueeze(0) - action_horizon = int(action.shape[1]) if action is not None else int((video0.shape[1] - 1) * dataset.action_video_freq_ratio) - infer_kwargs = { - "input_image": input_image, - "num_frames": int(video0.shape[1]), - "action_horizon": action_horizon, - "proprio": None if proprio is None else proprio[0, 0], - "num_inference_steps": self.eval_num_inference_steps, - "seed": sample_seed, - "rand_device": "cpu", - "tiled": self.eval_tiled, - } - if context is not None: - infer_kwargs.update(prompt=None, context=context[0], context_mask=context_mask[0]) - else: - infer_kwargs["prompt"] = prompt - prediction = model.infer(**infer_kwargs) - - pred_video_tensor = pil_frames_to_video_tensor(prediction["video"]) - gt_video_tensor = ((video0.detach().float().cpu().clamp(-1.0, 1.0) + 1.0) * 0.5).contiguous() - self._validate_video_shapes(pred_video_tensor, gt_video_tensor, "prediction", "ground truth") - - gt_video_batch = video0.unsqueeze(0).to(device=model.device, dtype=model.torch_dtype) - vae_latents = model._encode_video_latents(gt_video_batch, tiled=self.eval_tiled) - vae_video_tensor = pil_frames_to_video_tensor(model._decode_latents(vae_latents, tiled=self.eval_tiled)) - self._validate_video_shapes(vae_video_tensor, gt_video_tensor, "VAE reconstruction", "ground truth") - - pred_future = pred_video_tensor[:, 1:] - vae_future = vae_video_tensor[:, 1:] - gt_future = gt_video_tensor[:, 1:] - result.update( - { - "future_video_psnr_pred_gt": video_psnr(pred_future, gt_future), - "future_video_ssim_pred_gt": video_ssim(pred_future, gt_future), - "future_video_psnr_pred_vae": video_psnr(pred_future, vae_future), - "future_video_ssim_pred_vae": video_ssim(pred_future, vae_future), - "future_video_psnr_vae_gt": video_psnr(vae_future, gt_future), - "future_video_ssim_vae_gt": video_ssim(vae_future, gt_future), - } - ) - - pred_action = prediction.get("action") - if action is not None and pred_action is not None: - result.update(self._action_metrics(dataset, sample, pred_action, action)) - - if self.eval_save_video or self.eval_save_preview: - comparison_frames = self._comparison_frames(pred_video_tensor, vae_video_tensor, gt_video_tensor) - artifact_stem = f"sample_{sample_number:03d}_index_{eval_index:06d}" - if self.eval_save_video: - video_path = os.path.join(step_dir, f"{artifact_stem}.mp4") - save_mp4(comparison_frames, video_path, fps=self.eval_fps) - result["video_path"] = video_path - if self.eval_save_preview: - preview_path = os.path.join(step_dir, f"{artifact_stem}_preview.png") - comparison_frames[len(comparison_frames) // 2].save(preview_path) - result["preview_path"] = preview_path - return result - - def _action_metrics(self, dataset, sample, pred_action, target_action): - processor = dataset.processor - proprio = sample.get("proprio") - if proprio is None: - raise ValueError("Eval action denormalization requires sample['proprio'].") - - denormalized = {} - action_meta = processor.shape_meta["action"] - state_meta = processor.shape_meta["state"] - for name, raw_action in (("pred", pred_action), ("gt", target_action)): - action_btd = raw_action.unsqueeze(0) if raw_action.ndim == 2 else raw_action - action_btd = action_btd.detach().to(device="cpu", dtype=torch.float32) - batch = { - "action": action_btd, - "state": proprio.detach().to(device="cpu", dtype=torch.float32), - } - batch = processor.action_state_merger.backward(batch) - batch = processor.normalizer.backward(batch) - merged_batch = { - "action": {meta["key"]: batch["action"][meta["key"]].squeeze(0) for meta in action_meta}, - "state": {meta["key"]: batch["state"][meta["key"]].squeeze(0) for meta in state_meta}, - } - denormalized[name] = processor.action_state_merger.forward(merged_batch)["action"].float() - - diff = denormalized["pred"] - denormalized["gt"] - action_is_pad = sample.get("action_is_pad") - if action_is_pad is not None: - valid = ~action_is_pad[0].detach().cpu().bool() - if valid.any(): - diff = diff[valid] - return { - "action_l1": float(diff.abs().mean().item()), - "action_l2": float(diff.pow(2).mean().item()), - } - - @staticmethod - def _validate_video_shapes(left, right, left_name, right_name): - if left.shape != right.shape: - raise ValueError(f"Eval {left_name}/{right_name} shape mismatch: {tuple(left.shape)} vs {tuple(right.shape)}") - if left.shape[1] <= 1: - raise ValueError(f"Eval video must contain at least one future frame, got {tuple(left.shape)}") - - @staticmethod - def _comparison_frames(pred, vae, gt): - labels = ("Prediction", "VAE reconstruction", "Ground truth") - videos = (pred, vae, gt) - frame_height = int(gt.shape[2]) - frame_width = int(gt.shape[3]) - label_height = 32 - frames = [] - for frame_index in range(gt.shape[1]): - canvas = Image.new("RGB", (frame_width * 3, frame_height + label_height), color=(18, 18, 18)) - draw = ImageDraw.Draw(canvas) - for column, (label, video) in enumerate(zip(labels, videos)): - array = video[:, frame_index].permute(1, 2, 0).clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).numpy() - canvas.paste(Image.fromarray(np.ascontiguousarray(array)), (column * frame_width, label_height)) - draw.text((column * frame_width + 8, 9), label, fill=(245, 245, 245)) - frames.append(canvas) - return frames - - @staticmethod - def _gather_eval_results(local_results): - if not is_distributed(): - return local_results - gathered = [None for _ in range(get_world_size())] - torch.distributed.all_gather_object(gathered, local_results) - return sorted( - [result for rank_results in gathered for result in rank_results], - key=lambda result: result["sample_number"], - ) - - @staticmethod - def _summarize_eval_results(results): - if not results: - return {} - metric_keys = sorted(key for key, value in results[0].items() if isinstance(value, float)) - return {key: float(sum(result[key] for result in results if key in result) / sum(key in result for result in results)) for key in metric_keys} diff --git a/lightx2v_train/lightx2v_train/utils/registry.py b/lightx2v_train/lightx2v_train/utils/registry.py index a5047acec..a3c67aa94 100644 --- a/lightx2v_train/lightx2v_train/utils/registry.py +++ b/lightx2v_train/lightx2v_train/utils/registry.py @@ -89,7 +89,6 @@ def merge(self, other_register): "wan_t2v_ar": "lightx2v_train.model_zoo.wan.wan_t2v", "wan_t2v_14b": "lightx2v_train.model_zoo.wan.wan_t2v", "wan_t2v_14b_ar": "lightx2v_train.model_zoo.wan.wan_t2v", - "wan_fastwam": "lightx2v_train.model_zoo.wan.wan_fastwam", "wan_ti2v_5b": "lightx2v_train.model_zoo.wan.wan_ti2v_5b", "wan_ti2v_5b_ar": "lightx2v_train.model_zoo.wan.wan_ti2v_5b", } @@ -104,7 +103,6 @@ def merge(self, other_register): "sgmd": "lightx2v_train.trainers.sgmd", "teacher_forcing": "lightx2v_train.trainers.teacher_forcing", "training_cache": "lightx2v_train.trainers.training_cache", - "world_action": "lightx2v_train.trainers.world_action", } _INFERENCER_MODULES = { @@ -144,8 +142,6 @@ def _ensure_data_registered(data_name): return if data_name == "image_dataset": import lightx2v_train.data.image_dataset # noqa: F401 - elif data_name == "libero_fastwam_dataset": - import lightx2v_train.data.libero.dataset # noqa: F401 elif data_name in {"latent_dataset", "prompt_dataset", "video_dataset"}: import lightx2v_train.data.video_dataset # noqa: F401 diff --git a/lightx2v_train/scripts/run_fastwam_libero_plus_train_8gpu.sh b/lightx2v_train/scripts/run_fastwam_libero_plus_train_8gpu.sh deleted file mode 100755 index 0069aaf2d..000000000 --- a/lightx2v_train/scripts/run_fastwam_libero_plus_train_8gpu.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash - -export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7} - -NPROC_PER_NODE=${NPROC_PER_NODE:-8} - -torchrun \ ---standalone \ ---nproc_per_node="${NPROC_PER_NODE}" \ -train.py --config configs/train/fastwam/libero_plus_uncond_2cam224_8gpu.yaml diff --git a/lightx2v_train/scripts/run_fastwam_libero_train.sh b/lightx2v_train/scripts/run_fastwam_libero_train.sh deleted file mode 100755 index 6d93f44f3..000000000 --- a/lightx2v_train/scripts/run_fastwam_libero_train.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash - -export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3} - -NPROC_PER_NODE=${NPROC_PER_NODE:-4} - -torchrun \ ---standalone \ ---nproc_per_node="${NPROC_PER_NODE}" \ -train.py --config configs/train/fastwam/libero_uncond_2cam224.yaml diff --git a/lightx2v_train/train.py b/lightx2v_train/train.py index e201c8dea..193f62aa0 100644 --- a/lightx2v_train/train.py +++ b/lightx2v_train/train.py @@ -3,9 +3,14 @@ import torch from loguru import logger -from lightx2v_train.data import build_data, build_sample_processor, prepare_data +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.runtime import ( + cleanup_distributed, + init_distributed, + load_config, + setup_logger, +) from lightx2v_train.trainers import build_trainer @@ -24,7 +29,6 @@ def main(): setup_logger(config) try: - prepare_data(config) use_training_cache = config["data"].get("use_training_cache", False) sample_processor = None if use_training_cache else build_sample_processor(config) dataloader_train = build_data(config, train_or_val="train", sample_processor=sample_processor) From 96e2dea169a8e847aaf5d1e6f0b20ad10e70a40f Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 18:27:59 +0000 Subject: [PATCH 25/34] update cache --- lightx2v_train/cache_data.py | 11 ++++------- lightx2v_train/data_process/README.md | 4 ++-- .../lightx2v_train/trainers/training_cache.py | 3 +-- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/lightx2v_train/cache_data.py b/lightx2v_train/cache_data.py index 49a658072..6cddca83c 100644 --- a/lightx2v_train/cache_data.py +++ b/lightx2v_train/cache_data.py @@ -17,11 +17,9 @@ 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("--split", default="train") - parser.add_argument("--save-dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + 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("--max-samples", type=int) parser.add_argument("--overwrite", action="store_true") return parser.parse_args() @@ -33,10 +31,9 @@ def main(): "output_dir": args.output_dir, "save_dtype": args.save_dtype, "seed": args.seed, - "max_samples": args.max_samples, "overwrite": args.overwrite, } - config["data"][args.split]["preserve_records"] = True + config["data"]["train"]["preserve_records"] = True torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True @@ -45,7 +42,7 @@ def main(): try: sample_processor = build_sample_processor(config) - dataloader = build_data(config, train_or_val=args.split, sample_processor=sample_processor) + dataloader = build_data(config, train_or_val="train", sample_processor=sample_processor) model = build_model(config) model.load_components( diff --git a/lightx2v_train/data_process/README.md b/lightx2v_train/data_process/README.md index f70d93246..d9ed8f0e1 100644 --- a/lightx2v_train/data_process/README.md +++ b/lightx2v_train/data_process/README.md @@ -18,8 +18,8 @@ cd /path/to/LightX2V python lightx2v_train/cache_data.py \ --config lightx2v_train/configs/train/flow/qwen_image_lora.yaml \ - --output-dir /path/to/qwen_image_cache \ - --save-dtype bf16 + --output_dir /path/to/qwen_image_cache \ + --save_dtype bf16 ``` 输出格式对所有图像模型保持一致： diff --git a/lightx2v_train/lightx2v_train/trainers/training_cache.py b/lightx2v_train/lightx2v_train/trainers/training_cache.py index f80dd4e29..966bb8197 100644 --- a/lightx2v_train/lightx2v_train/trainers/training_cache.py +++ b/lightx2v_train/lightx2v_train/trainers/training_cache.py @@ -121,8 +121,7 @@ def train(self): cache_dir.mkdir(parents=True, exist_ok=True) dataset = self.dataloader.dataset - max_samples = self.cache_config["max_samples"] - sample_count = len(dataset) if max_samples is None else min(len(dataset), max_samples) + sample_count = len(dataset) dtype = CACHE_DTYPES[self.cache_config["save_dtype"]] records = [] From d3f20a765784d9ed2d9472a334845ac384df2ea3 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 19:29:32 +0000 Subject: [PATCH 26/34] update cache config --- .../train/flow/qwen_image_lora_cache.yaml | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 lightx2v_train/configs/train/flow/qwen_image_lora_cache.yaml 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 From 8e0551aac1ca917286adb2aab8b851194b93a862 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 20:01:12 +0000 Subject: [PATCH 27/34] fix --- lightx2v_train/lightx2v_train/infer/image.py | 14 +- .../model_zoo/flux2/data_process.py | 20 ++- .../model_zoo/flux2/flux2_klein.py | 1 + .../model_zoo/qwen_image/data_process.py | 19 ++- .../model_zoo/qwen_image/qwen_image.py | 1 + .../trainers/consistency/pcm.py | 20 ++- .../tests/test_image_training_regressions.py | 126 ++++++++++++++++++ 7 files changed, 189 insertions(+), 12 deletions(-) create mode 100644 lightx2v_train/tests/test_image_training_regressions.py diff --git a/lightx2v_train/lightx2v_train/infer/image.py b/lightx2v_train/lightx2v_train/infer/image.py index b224f34c0..eb75483ff 100644 --- a/lightx2v_train/lightx2v_train/infer/image.py +++ b/lightx2v_train/lightx2v_train/infer/image.py @@ -15,7 +15,7 @@ def _record_has_source_images(record): @INFERENCER_REGISTER("image_infer") class ImageInferencer(BaseInferencer): - def _inference_sigmas(self, num_inference_steps): + 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: @@ -25,7 +25,12 @@ def _inference_sigmas(self, num_inference_steps): from lightx2v_train.trainers.consistency.pcm import pcm_inference_sigmas - return pcm_inference_sigmas(num_inference_steps, int(pcm_solver_steps)) + 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] @@ -107,7 +112,10 @@ def infer(self): latent_hw = (latent.shape[-2], latent.shape[-1]) self.scheduler.set_timesteps( num_inference_steps, - sigmas=self._inference_sigmas(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) diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py index c88292c54..ba3b825eb 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/data_process.py @@ -95,6 +95,7 @@ def infer_target_size(self, sample, default_height, default_width): 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: @@ -108,15 +109,28 @@ def __call__(self, sample): source_sizes.append(size) reference_height, reference_width = source_sizes[0] - sample["meta"]["reference_image_height"] = reference_height - sample["meta"]["reference_image_width"] = reference_width + 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: - inputs["target_pixel_values"] = self._process_target( + 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): diff --git a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py index ed7289a17..90079311c 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py +++ b/lightx2v_train/lightx2v_train/model_zoo/flux2/flux2_klein.py @@ -38,6 +38,7 @@ def register_capabilities(self): channels_path="transformer.config.in_channels", spatial_downsample_multiplier=2, ), + guidance_in_denoiser_space=True, ), ) self.capabilities.register(ConsistencyModelCapability, Flux2ConsistencyModelCapability(self)) diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py index 5966bb40d..946777603 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/data_process.py @@ -109,20 +109,33 @@ def _resolve_target_size(self, sample, ratio, target_area, area_multiple=None): class QwenImageEditDataProcessor(QwenImageDataProcessor): def __call__(self, sample): inputs = sample["inputs"] + meta = sample["meta"] target_image = inputs.pop("target_image", None) source_images = inputs.pop("source_images", []) reference_image = source_images[-1] if source_images else None if reference_image is not None: - sample["meta"]["reference_image_height"] = reference_image.height - sample["meta"]["reference_image_width"] = reference_image.width + meta["reference_image_height"] = reference_image.height + meta["reference_image_width"] = reference_image.width if target_image is not None: - inputs["target_pixel_values"] = self._process_target( + target_pixels = self._process_target( target_image, sample, reference_image=reference_image, area_multiple=32, ) + inputs["target_pixel_values"] = target_pixels + meta["target_height"] = int(target_pixels.shape[-2]) + meta["target_width"] = int(target_pixels.shape[-1]) + elif reference_image is not None: + target_width, target_height = self._resolve_target_size( + sample, + reference_image.width / reference_image.height, + self.target_area, + area_multiple=32, + ) + meta["target_height"] = target_height + meta["target_width"] = target_width if source_images: condition_images, vae_images = self._process_source_images(source_images) inputs["source_condition_images"] = condition_images diff --git a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py index 22b43a4cf..933c0066c 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py +++ b/lightx2v_train/lightx2v_train/model_zoo/qwen_image/qwen_image.py @@ -46,6 +46,7 @@ def register_capabilities(self): channels_path="latent_channels", temporal_size=1, ), + guidance_in_denoiser_space=True, ), ) self.capabilities.register( diff --git a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py index 345192fe4..0f653fe8c 100644 --- a/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -99,8 +99,19 @@ def __call__(self, prediction: Tensor, target: Tensor) -> Tensor: return elementwise.flatten(1).mean(dim=1) -def pcm_inference_sigmas(num_inference_steps: int, num_solver_steps: int) -> list[float]: - """Return the unshifted descending phase schedule used by PCM inference.""" +def pcm_inference_sigmas( + num_inference_steps: int, + num_solver_steps: int, + *, + scheduler=None, + latent_hw=None, +) -> list[float]: + """Return the descending PCM phase schedule used by inference. + + PCM is trained against phase boundaries after the scheduler's time shift. + Custom inference sigmas otherwise bypass scheduler shifting, so apply the + same transform here when a scheduler is provided. + """ if not 1 <= num_inference_steps <= num_solver_steps: raise ValueError("PCM inference steps must be in [1, num_solver_steps].") indices = torch.div( @@ -108,7 +119,10 @@ def pcm_inference_sigmas(num_inference_steps: int, num_solver_steps: int) -> lis num_inference_steps, rounding_mode="floor", ) - return ((num_solver_steps - indices).float() / num_solver_steps).tolist() + sigmas = (num_solver_steps - indices).float() / num_solver_steps + if scheduler is not None and scheduler.do_time_shift: + sigmas = scheduler.time_shift(sigmas, latent_hw=latent_hw) + return sigmas.tolist() @CONSISTENCY_OBJECTIVE_REGISTER("pcm") diff --git a/lightx2v_train/tests/test_image_training_regressions.py b/lightx2v_train/tests/test_image_training_regressions.py new file mode 100644 index 000000000..2d25319de --- /dev/null +++ b/lightx2v_train/tests/test_image_training_regressions.py @@ -0,0 +1,126 @@ +import torch +from PIL import Image + +from lightx2v_train.infer.image import ImageInferencer +from lightx2v_train.model_capabilities import DistributionMatchingCapability +from lightx2v_train.model_zoo.flux2.data_process import Flux2EditDataProcessor +from lightx2v_train.model_zoo.flux2.flux2_klein import Flux2KleinModel +from lightx2v_train.model_zoo.flux2.flux2_klein_edit import Flux2KleinEditModel +from lightx2v_train.model_zoo.qwen_image.data_process import QwenImageEditDataProcessor +from lightx2v_train.model_zoo.qwen_image.qwen_image import QwenImageModel +from lightx2v_train.model_zoo.qwen_image.qwen_image_edit import QwenImageEditModel + + +def _processor_config(target_area=64 * 96): + return { + "model": { + "running_dtype": "fp32", + "input_preprocessing": {"target_area": target_area}, + }, + "data": { + "processor": { + "size_multiple": 16, + "target_area": target_area, + }, + }, + } + + +def _edit_sample(*, target=True, target_size=None): + inputs = { + "source_images": [Image.new("RGB", (96, 64))], + } + if target: + inputs["target_image"] = Image.new("RGB", (150, 100)) + meta = {} + if target_size is not None: + meta["target_height"], meta["target_width"] = target_size + return { + "inputs": inputs, + "conditioning": {"prompt": "edit"}, + "meta": meta, + } + + +def test_pcm_inference_applies_the_training_time_shift(): + config = { + "model": {"running_dtype": "fp32"}, + "scheduler": { + "time_shift_settings": { + "do_time_shift": True, + "dynamic_shift": True, + "shift_x1": 256, + "shift_x2": 4096, + "shift_y1": 0.5, + "shift_y2": 1.15, + "patch_size": [2, 2], + }, + }, + "inference": {"pcm_solver_steps": 100}, + } + inferencer = ImageInferencer(config) + latent_hw = (128, 128) + actual = torch.tensor( + inferencer._inference_sigmas(4, latent_hw=latent_hw), + dtype=torch.float32, + ) + raw_phase_boundaries = torch.tensor([1.0, 0.75, 0.5, 0.25]) + expected = inferencer.scheduler.time_shift( + raw_phase_boundaries, + latent_hw=latent_hw, + ) + + torch.testing.assert_close(actual, expected) + assert not torch.allclose(actual, raw_phase_boundaries) + + +def test_qwen_edit_records_actual_dmd_target_size(): + processor = QwenImageEditDataProcessor(_processor_config()) + processor._process_source_images = lambda images: ([], []) + + with_target = processor(_edit_sample(target=True)) + pixels = with_target["inputs"]["target_pixel_values"] + assert with_target["meta"]["target_height"] == pixels.shape[-2] + assert with_target["meta"]["target_width"] == pixels.shape[-1] + + without_target = processor(_edit_sample(target=False)) + assert without_target["meta"]["target_height"] > 0 + assert without_target["meta"]["target_width"] > 0 + assert without_target["meta"]["target_height"] % 32 == 0 + assert without_target["meta"]["target_width"] % 32 == 0 + + +def test_flux2_edit_records_actual_dmd_target_size(): + processor = Flux2EditDataProcessor(_processor_config()) + processor._process_source = lambda image: ( + torch.zeros(3, 64, 96), + (64, 96), + ) + + with_target = processor(_edit_sample(target=True)) + pixels = with_target["inputs"]["target_pixel_values"] + assert with_target["meta"]["target_height"] == pixels.shape[-2] + assert with_target["meta"]["target_width"] == pixels.shape[-1] + + without_target = processor( + _edit_sample(target=False, target_size=(70, 100)), + ) + assert without_target["meta"]["target_height"] % 16 == 0 + assert without_target["meta"]["target_width"] % 16 == 0 + + +def test_qwen_and_flux2_klein_apply_dmd_cfg_norm_to_packed_tokens(): + config = {"model": {"running_dtype": "fp32"}} + model_classes = ( + QwenImageModel, + QwenImageEditModel, + Flux2KleinModel, + Flux2KleinEditModel, + ) + + for model_class in model_classes: + model = model_class(config) + capability = model.ensure_capabilities().require( + DistributionMatchingCapability, + ) + assert capability._guidance_in_denoiser_space is True From 678f9724bca42125d0c95145216ba05bd072ecff Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Mon, 24 Aug 2026 20:41:46 +0000 Subject: [PATCH 28/34] remove ltx --- .../train/dmd/ltx_t2av_ar_dmd_lora.yaml | 127 -- .../configs/train/dmd/ltx_t2av_dmd_lora.yaml | 175 --- .../configs/train/flow/ltx_t2av_full.yaml | 75 -- .../configs/train/flow/ltx_t2av_lora.yaml | 107 -- .../configs/train/flow/ltx_t2av_train.yaml | 107 -- .../tf/ltx_t2av_teacher_forcing_lmdb.yaml | 93 -- lightx2v_train/data_process/README.md | 54 +- .../data_process/ltx/build_latent_dataset.py | 198 --- .../lightx2v_train/model_zoo/__init__.py | 2 - .../lightx2v_train/model_zoo/ltx/__init__.py | 1 - .../ltx/capability_adapters/__init__.py | 5 - .../ltx/capability_adapters/common.py | 122 -- ...essive_distribution_matching_capability.py | 387 ------ .../ltx_consistency_model_capability.py | 194 --- .../ltx_distribution_matching_capability.py | 438 ------- .../ltx_flow_matching_capability.py | 143 --- .../ltx_teacher_forcing_capability.py | 357 ------ .../lightx2v_train/model_zoo/ltx/ltx_t2av.py | 416 ------- .../model_zoo/native/ltx2/__init__.py | 68 - .../native/ltx2/ltx_core/__init__.py | 0 .../native/ltx2/ltx_core/batch_split.py | 91 -- .../ltx2/ltx_core/block_streaming/__init__.py | 20 - .../ltx_core/block_streaming/block_fetcher.py | 82 -- .../ltx2/ltx_core/block_streaming/builder.py | 463 ------- .../ltx2/ltx_core/block_streaming/disk.py | 161 --- .../ltx2/ltx_core/block_streaming/pool.py | 69 -- .../ltx2/ltx_core/block_streaming/provider.py | 138 --- .../ltx2/ltx_core/block_streaming/source.py | 171 --- .../ltx2/ltx_core/block_streaming/utils.py | 165 --- .../ltx2/ltx_core/block_streaming/wrapper.py | 95 -- .../ltx2/ltx_core/components/__init__.py | 10 - .../ltx_core/components/diffusion_steps.py | 183 --- .../ltx2/ltx_core/components/guiders.py | 361 ------ .../ltx2/ltx_core/components/noisers.py | 36 - .../ltx2/ltx_core/components/patchifiers.py | 352 ------ .../ltx2/ltx_core/components/protocols.py | 98 -- .../ltx2/ltx_core/components/schedulers.py | 125 -- .../ltx2/ltx_core/conditioning/__init__.py | 23 - .../ltx2/ltx_core/conditioning/exceptions.py | 4 - .../native/ltx2/ltx_core/conditioning/item.py | 20 - .../ltx2/ltx_core/conditioning/mask_utils.py | 200 --- .../ltx_core/conditioning/types/__init__.py | 17 - .../types/attention_strength_wrapper.py | 70 -- .../conditioning/types/keyframe_cond.py | 83 -- .../conditioning/types/latent_cond.py | 40 - .../ltx_core/conditioning/types/mask_cond.py | 48 - .../conditioning/types/noise_mask_cond.py | 43 - .../types/reference_audio_cond.py | 58 - .../types/reference_video_cond.py | 101 -- .../native/ltx2/ltx_core/guidance/__init__.py | 15 - .../ltx2/ltx_core/guidance/perturbations.py | 77 -- .../model_zoo/native/ltx2/ltx_core/hdr.py | 71 -- .../native/ltx2/ltx_core/loader/__init__.py | 56 - .../ltx2/ltx_core/loader/attention_ops.py | 39 - .../native/ltx2/ltx_core/loader/fuse_loras.py | 177 --- .../native/ltx2/ltx_core/loader/helpers.py | 60 - .../native/ltx2/ltx_core/loader/kernels.py | 79 -- .../native/ltx2/ltx_core/loader/module_ops.py | 14 - .../native/ltx2/ltx_core/loader/primitives.py | 164 --- .../native/ltx2/ltx_core/loader/registry.py | 84 -- .../native/ltx2/ltx_core/loader/sd_ops.py | 132 -- .../native/ltx2/ltx_core/loader/sft_loader.py | 65 - .../loader/single_gpu_model_builder.py | 231 ---- .../native/ltx2/ltx_core/modality_tiling.py | 227 ---- .../native/ltx2/ltx_core/model/__init__.py | 8 - .../ltx2/ltx_core/model/audio_vae/__init__.py | 29 - .../ltx_core/model/audio_vae/attention.py | 70 -- .../ltx_core/model/audio_vae/audio_vae.py | 501 -------- .../model/audio_vae/causal_conv_2d.py | 109 -- .../model/audio_vae/causality_axis.py | 10 - .../ltx_core/model/audio_vae/downsample.py | 109 -- .../model/audio_vae/model_configurator.py | 193 --- .../ltx2/ltx_core/model/audio_vae/ops.py | 72 -- .../ltx2/ltx_core/model/audio_vae/resnet.py | 171 --- .../ltx2/ltx_core/model/audio_vae/upsample.py | 105 -- .../ltx2/ltx_core/model/audio_vae/vocoder.py | 591 --------- .../ltx2/ltx_core/model/common/__init__.py | 9 - .../ltx_core/model/common/normalization.py | 57 - .../ltx2/ltx_core/model/model_protocol.py | 39 - .../ltx_core/model/transformer/__init__.py | 22 - .../ltx2/ltx_core/model/transformer/adaln.py | 44 - .../ltx_core/model/transformer/attention.py | 836 ------------- .../ltx_core/model/transformer/compiling.py | 126 -- .../model/transformer/feed_forward.py | 14 - .../ltx_core/model/transformer/gelu_approx.py | 10 - .../ltx_core/model/transformer/modality.py | 68 - .../ltx2/ltx_core/model/transformer/model.py | 769 ------------ .../model/transformer/model_configurator.py | 213 ---- .../ltx2/ltx_core/model/transformer/ops.py | 105 -- .../ltx2/ltx_core/model/transformer/rope.py | 204 --- .../model/transformer/text_projection.py | 38 - .../model/transformer/timestep_embedding.py | 143 --- .../ltx_core/model/transformer/transformer.py | 444 ------- .../model/transformer/transformer_args.py | 351 ------ .../ltx2/ltx_core/model/upsampler/__init__.py | 10 - .../model/upsampler/blur_downsample.py | 53 - .../ltx2/ltx_core/model/upsampler/model.py | 139 --- .../model/upsampler/model_configurator.py | 30 - .../ltx_core/model/upsampler/pixel_shuffle.py | 54 - .../ltx_core/model/upsampler/res_block.py | 37 - .../upsampler/spatial_rational_resampler.py | 46 - .../ltx2/ltx_core/model/video_vae/__init__.py | 25 - .../ltx_core/model/video_vae/convolution.py | 316 ----- .../ltx2/ltx_core/model/video_vae/enums.py | 20 - .../video_vae/memory_efficient_decode.py | 649 ---------- .../model/video_vae/model_configurator.py | 79 -- .../ltx_core/model/video_vae/normalization.py | 3 - .../ltx2/ltx_core/model/video_vae/ops.py | 78 -- .../ltx2/ltx_core/model/video_vae/resnet.py | 264 ---- .../ltx2/ltx_core/model/video_vae/sampling.py | 123 -- .../ltx2/ltx_core/model/video_vae/tiling.py | 65 - .../ltx_core/model/video_vae/video_vae.py | 1093 ----------------- .../ltx2/ltx_core/quantization/__init__.py | 17 - .../ltx2/ltx_core/quantization/fp8_cast.py | 322 ----- .../ltx_core/quantization/fp8_scaled_mm.py | 209 ---- .../ltx2/ltx_core/quantization/policy.py | 24 - .../quantization/trtllm_scaled_usable.py | 37 - .../ltx2/ltx_core/text_encoders/__init__.py | 1 - .../ltx_core/text_encoders/gemma/__init__.py | 33 - .../ltx_core/text_encoders/gemma/config.py | 75 -- .../gemma/embeddings_connector.py | 241 ---- .../gemma/embeddings_processor.py | 113 -- .../gemma/encoders/base_encoder.py | 211 ---- .../gemma/encoders/encoder_configurator.py | 254 ---- .../prompts/gemma_i2v_system_prompt.txt | 30 - .../prompts/gemma_t2v_system_prompt.txt | 40 - .../text_encoders/gemma/feature_extractor.py | 129 -- .../ltx_core/text_encoders/gemma/tokenizer.py | 67 - .../model_zoo/native/ltx2/ltx_core/tiling.py | 443 ------- .../model_zoo/native/ltx2/ltx_core/tools.py | 176 --- .../model_zoo/native/ltx2/ltx_core/types.py | 211 ---- .../model_zoo/native/ltx2/ltx_core/utils.py | 62 - .../lightx2v_train/utils/constants.py | 14 - .../lightx2v_train/utils/registry.py | 2 - lightx2v_train/scripts/run_ltx_t2av_train.sh | 9 - 135 files changed, 8 insertions(+), 19143 deletions(-) delete mode 100644 lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml delete mode 100644 lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml delete mode 100644 lightx2v_train/configs/train/flow/ltx_t2av_full.yaml delete mode 100644 lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml delete mode 100644 lightx2v_train/configs/train/flow/ltx_t2av_train.yaml delete mode 100644 lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml delete mode 100644 lightx2v_train/data_process/ltx/build_latent_dataset.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/batch_split.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/block_fetcher.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/builder.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/disk.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/pool.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/provider.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/source.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/utils.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/block_streaming/wrapper.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/diffusion_steps.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/guiders.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/noisers.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/patchifiers.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/protocols.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/components/schedulers.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/exceptions.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/item.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/mask_utils.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/attention_strength_wrapper.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/keyframe_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/latent_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/mask_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/noise_mask_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_audio_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/conditioning/types/reference_video_cond.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/guidance/perturbations.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/hdr.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/attention_ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/fuse_loras.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/helpers.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/kernels.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/module_ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/primitives.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/registry.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sd_ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/sft_loader.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/loader/single_gpu_model_builder.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/modality_tiling.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/attention.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/audio_vae.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causal_conv_2d.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/causality_axis.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/downsample.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/model_configurator.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/resnet.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/upsample.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/audio_vae/vocoder.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/common/normalization.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/model_protocol.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/adaln.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/attention.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/compiling.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/feed_forward.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/gelu_approx.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/modality.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/model_configurator.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/rope.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/text_projection.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/timestep_embedding.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/transformer/transformer_args.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/blur_downsample.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/model_configurator.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/pixel_shuffle.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/res_block.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/upsampler/spatial_rational_resampler.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/convolution.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/enums.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/memory_efficient_decode.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/model_configurator.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/normalization.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/ops.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/resnet.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/sampling.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/tiling.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/model/video_vae/video_vae.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_cast.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/fp8_scaled_mm.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/policy.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/quantization/trtllm_scaled_usable.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/__init__.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/config.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_connector.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/embeddings_processor.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/base_encoder.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/feature_extractor.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/text_encoders/gemma/tokenizer.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/tiling.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/tools.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/types.py delete mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/ltx2/ltx_core/utils.py delete mode 100644 lightx2v_train/scripts/run_ltx_t2av_train.sh 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 95bdab321..000000000 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_ar_dmd_lora.yaml +++ /dev/null @@ -1,127 +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 - 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: autoregressive_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 6bba4edf4..000000000 --- a/lightx2v_train/configs/train/dmd/ltx_t2av_dmd_lora.yaml +++ /dev/null @@ -1,175 +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 - 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: 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/flow/ltx_t2av_full.yaml b/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml deleted file mode 100644 index 2026108c4..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_full.yaml +++ /dev/null @@ -1,75 +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 - 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: flow_matching - 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 7858854b7..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_lora.yaml +++ /dev/null @@ -1,107 +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 - 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: 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 - 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 9936a3d6b..000000000 --- a/lightx2v_train/configs/train/flow/ltx_t2av_train.yaml +++ /dev/null @@ -1,107 +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 - 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: flow_matching - 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/tf/ltx_t2av_teacher_forcing_lmdb.yaml b/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml deleted file mode 100644 index 644522f65..000000000 --- a/lightx2v_train/configs/train/tf/ltx_t2av_teacher_forcing_lmdb.yaml +++ /dev/null @@ -1,93 +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 - 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: 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/data_process/README.md b/lightx2v_train/data_process/README.md index d9ed8f0e1..559faff5f 100644 --- a/lightx2v_train/data_process/README.md +++ b/lightx2v_train/data_process/README.md @@ -1,11 +1,10 @@ # 训练数据与预处理 -本文说明训练数据格式，以及 WAN 和 LTX2 的 latent 数据集构建方法。预处理结果包含统一的 `metadata.jsonl`，可由 `latent_dataset` 直接读取。 +本文说明训练数据格式，以及 WAN 的 latent 数据集构建方法。预处理结果包含统一的 `metadata.jsonl`，可由 `latent_dataset` 直接读取。 ```text data_process/ -├── wan/build_latent_dataset.py -└── ltx/build_latent_dataset.py +└── wan/build_latent_dataset.py ``` ## 图像训练缓存 @@ -73,12 +72,12 @@ worker、持久 worker 与预取用于隐藏单样本 PT 的读取开销。 | --- | --- | --- | | 图像 DMD | `prompt_dataset` | prompt、目标高度和目标宽度 | | Flow | `video_dataset` | 视频路径和 caption 元数据，训练时在线编码 | -| Flow | `latent_dataset` | 预先生成的视频、音频和文本缓存 | +| Flow | `latent_dataset` | 预先生成的视频和文本缓存 | | DMD / AR-DMD | `prompt_dataset` | TXT prompt，训练时在线编码 | | DMD / AR-DMD | `latent_dataset` | 预先生成的文本条件缓存 | | Teacher Forcing | `latent_dataset` | PT 缓存元数据或 LMDB | -WAN 和 LTX2 使用相同的数据集入口。所有数据集均返回统一结构： +所有数据集均返回统一结构： ```python { @@ -195,7 +194,7 @@ rank 数整除；设置 `drop_last: true` 则丢弃不足一个分布式训练 | 字段 | 含义 | | --- | --- | | `video_latent_path` | 视频 latent 的 PT 文件 | -| `audio_latent_path` | 音频 latent 的 PT 文件，LTX2 T2AV 使用 | +| `audio_latent_path` | 音频 latent 的 PT 文件，音视频模型可选 | | `condition_path` | 正向文本条件的 PT 文件 | | `negative_condition_path` | 当前样本的负向文本条件，可选 | | `caption` | 原始文本，可选 | @@ -206,12 +205,6 @@ WAN 示例： {"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: @@ -283,40 +276,9 @@ wan_latent_cache/ 脚本支持连续片段拼接。文件名形如 `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 的写法一致。 +训练任务通过配置中的 `data.train` 选择数据集。 ### Flow @@ -349,7 +311,7 @@ data: shuffle: true ``` -WAN 的缓存记录包含 `video_latent_path` 和 `condition_path`。LTX2 T2AV 还需要 `audio_latent_path`。 +WAN 的缓存记录包含 `video_latent_path` 和 `condition_path`。 ### DMD / AR-DMD @@ -405,4 +367,4 @@ data: drop_last: true ``` -WAN Teacher Forcing 需要视频 latent 和文本条件。LTX2 T2AV Teacher Forcing 需要视频 latent、音频 latent 和文本条件。 +WAN Teacher Forcing 需要视频 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/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index faf1ee464..b9984b3a0 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -10,7 +10,6 @@ "LingBotVideoModel": (".wan.lingbot_video", "LingBotVideoModel"), "LongCatImageModel": (".longcat_image.longcat_image", "LongCatImageModel"), "LongCatImageEditModel": (".longcat_image.longcat_image_edit", "LongCatImageEditModel"), - "LTX2T2AVModel": (".ltx.ltx_t2av", "LTX2T2AVModel"), "MiniMaxH3T2AVModel": (".minimax_h3.minimax_h3_t2av", "MiniMaxH3T2AVModel"), "QwenImageModel": (".qwen_image.qwen_image", "QwenImageModel"), "QwenImageEditModel": (".qwen_image.qwen_image_edit", "QwenImageEditModel"), @@ -59,7 +58,6 @@ def __getattr__(name): "Flux2KleinModel", "Flux2KleinEditModel", "LingBotVideoModel", - "LTX2T2AVModel", "MiniMaxH3T2AVModel", "WanT2VModel", "WanTI2V5BModel", diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py deleted file mode 100644 index c58c1964d..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""LTX model wrappers.""" diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py deleted file mode 100644 index 3ad2b5540..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Capabilities implemented specifically by LTX models.""" - -from .ltx_consistency_model_capability import LTXConsistencyModelCapability - -__all__ = ["LTXConsistencyModelCapability"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py deleted file mode 100644 index 5d1cd3784..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/common.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Shared latent value objects for LTX training capabilities.""" - -from dataclasses import dataclass - -import torch -from torch import Tensor - -from lightx2v_train.model_zoo.native.ltx2 import ( - AudioLatentShape, - AudioPatchifier, - SpatioTemporalScaleFactors, - VideoLatentPatchifier, - VideoLatentShape, - get_pixel_coords, -) - - -@dataclass(frozen=True) -class LTXJointLatentShape: - video_latent: tuple[int, ...] - audio_latent: tuple[int, ...] - video_tokens: tuple[int, ...] - audio_tokens: tuple[int, ...] - fps: float - - -@dataclass(frozen=True) -class LTXJointLatents: - video: Tensor - audio: Tensor - shape: LTXJointLatentShape - - -class LTXLatentCodec: - """Validate, patchify, position, and flatten joint LTX latent tensors.""" - - def __init__(self, model, *, default_fps: float = 24.0) -> None: - self.model = model - self.default_fps = float(default_fps) - self.video_patchifier = VideoLatentPatchifier(patch_size=1) - self.audio_patchifier = AudioPatchifier(patch_size=1) - - def prepare(self, inputs, dtype): - video_data = inputs.get("video_latents") - audio_data = inputs.get("audio_latents") - if video_data is None or audio_data is None: - raise KeyError("LTX training expects inputs.video_latents and inputs.audio_latents.") - video_latents = self.video_latents(video_data, dtype) - audio_latents = self.audio_latents(audio_data, dtype) - video_tokens = self.video_patchifier.patchify(video_latents) - audio_tokens = self.audio_patchifier.patchify(audio_latents) - return video_data, audio_data, video_latents, audio_latents, video_tokens, audio_tokens - - def positions(self, video_data, video_latents, audio_tokens): - video_positions = self.video_positions(video_data, video_latents) - audio_positions = self.audio_patchifier.get_patch_grid_bounds( - output_shape=AudioLatentShape( - frames=audio_tokens.shape[1], - mel_bins=16, - batch=1, - channels=8, - ), - device=audio_tokens.device, - ) - return video_positions, audio_positions - - def flatten(self, video_tokens: Tensor, audio_tokens: Tensor) -> Tensor: - if video_tokens.shape[0] != 1 or audio_tokens.shape[0] != 1: - raise ValueError("LTX consistency training only supports physical batch size 1.") - return torch.cat((video_tokens.flatten(1), audio_tokens.flatten(1)), dim=1) - - @staticmethod - def unflatten(flat: Tensor, video_shape, audio_shape) -> tuple[Tensor, Tensor]: - video_numel = int(torch.Size(video_shape).numel()) - audio_numel = int(torch.Size(audio_shape).numel()) - if flat.shape != (1, video_numel + audio_numel): - raise ValueError(f"Flattened LTX latent must have shape {(1, video_numel + audio_numel)}, got {tuple(flat.shape)}.") - video, audio = flat.split((video_numel, audio_numel), dim=1) - return video.reshape(video_shape), audio.reshape(audio_shape) - - def video_latents(self, data, dtype): - latents = data["latents"].to(device=self.model.device, dtype=dtype) - if latents.ndim == 4: - latents = latents.unsqueeze(0) - if latents.ndim != 5: - raise ValueError(f"LTX video latents must have shape [B,C,F,H,W] or [C,F,H,W], got {tuple(latents.shape)}.") - if latents.shape[0] != 1: - raise ValueError("LTX training only supports physical batch size 1.") - return latents.contiguous() - - def audio_latents(self, data, dtype): - latents = data["latents"].to(device=self.model.device, dtype=dtype) - if latents.ndim == 3: - latents = latents.unsqueeze(0) - if latents.ndim != 4: - raise ValueError(f"LTX audio latents must have shape [B,C,T,F] or [C,T,F], got {tuple(latents.shape)}.") - if latents.shape[0] != 1: - raise ValueError("LTX training only supports physical batch size 1.") - return latents.contiguous() - - def video_positions(self, data, latents): - frames = int(data["num_frames"][0].item()) if "num_frames" in data else latents.shape[2] - height = int(data["height"][0].item()) if "height" in data else latents.shape[3] - width = int(data["width"][0].item()) if "width" in data else latents.shape[4] - fps = float(data["fps"][0].item()) if "fps" in data else self.default_fps - latent_coords = self.video_patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape( - frames=frames, - height=height, - width=width, - batch=1, - channels=128, - ), - device=latents.device, - ) - positions = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=SpatioTemporalScaleFactors.default(), - causal_fix=True, - ).float() - positions[:, 0, ...] /= fps - return positions diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py deleted file mode 100644 index ca85b2d59..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_autoregressive_distribution_matching_capability.py +++ /dev/null @@ -1,387 +0,0 @@ -"""Autoregressive distribution-matching capability for LTX models.""" - -from __future__ import annotations - -import math - -import torch - -from lightx2v_train.model_capabilities import ( - AutoregressiveDistributionMatchingCapability, - AutoregressiveRolloutContext, - BoundCapability, -) -from lightx2v_train.model_zoo.native.ltx2 import Modality - -from .common import LTXJointLatents -from .ltx_distribution_matching_capability import LTXDistributionMatchingCapability - - -class LTXAutoregressiveDistributionMatchingCapability( - BoundCapability, - AutoregressiveDistributionMatchingCapability, -): - """Cached chunk-wise rollout for LTX joint video/audio states.""" - - def __init__(self, model) -> None: - super().__init__(model) - self.distribution_matching = LTXDistributionMatchingCapability(model) - - def rollout( - self, - condition, - latent_shape, - initial_latents, - context: AutoregressiveRolloutContext, - ): - latents = initial_latents - if latents.video.shape[0] != 1 or latents.audio.shape[0] != 1: - raise ValueError("LTX autoregressive DMD only supports physical batch size 1.") - _, _, num_frames, latent_height, latent_width = latent_shape.video_latent - frame_sequence_length = int(latent_height) * int(latent_width) - if latents.video.shape[1] != num_frames * frame_sequence_length: - if latents.video.shape[1] % num_frames: - raise ValueError("LTX video token count is not divisible by latent frames.") - frame_sequence_length = latents.video.shape[1] // num_frames - if num_frames % context.frames_per_chunk: - raise ValueError(f"LTX latent frames ({num_frames}) must be divisible by frames_per_chunk ({context.frames_per_chunk}).") - - transformer = self._unwrap_transformer() - if transformer is None: - raise RuntimeError("LTX autoregressive rollout requires transformer KV-cache inference support.") - transformer.train() - num_blocks = num_frames // context.frames_per_chunk - exit_indices = self._sample_exit_indices( - 1 if context.same_step_across_blocks else num_blocks, - len(context.denoising_steps), - latents.video.device, - ) - video_positions = self.distribution_matching._video_token_positions(latent_shape) - audio_positions = self.distribution_matching._audio_token_positions(latent_shape) - cache = self._new_caches( - dtype=latents.video.dtype, - device=latents.video.device, - video_total_tokens=latents.video.shape[1], - audio_total_tokens=latents.audio.shape[1], - video_frame_sequence_length=frame_sequence_length, - num_frames=num_frames, - frames_per_chunk=context.frames_per_chunk, - ) - - video_chunks = [] - audio_chunks = [] - for block_index in range(num_blocks): - frame_start = block_index * context.frames_per_chunk - frame_end = frame_start + context.frames_per_chunk - video_start = frame_start * frame_sequence_length - video_end = frame_end * frame_sequence_length - audio_start, audio_end = self._audio_range( - block_index, - num_blocks, - latents.audio.shape[1], - ) - block = LTXJointLatents( - latents.video[:, video_start:video_end], - latents.audio[:, audio_start:audio_end], - latent_shape, - ) - positions = ( - video_positions[:, :, video_start:video_end], - audio_positions[:, :, audio_start:audio_end], - ) - exit_index = int(exit_indices[0] if context.same_step_across_blocks else exit_indices[block_index]) - - x0 = None - for step_index in range(exit_index + 1): - sigma = context.trajectory_scheduler.sigma_at( - step_index, - device=self.model.device, - dtype=context.running_dtype, - ) - enable_grad = context.grad_enabled and step_index == exit_index - grad_context = torch.enable_grad if enable_grad else torch.no_grad - with grad_context(): - velocity = self._predict_chunk( - block, - sigma, - condition, - positions, - cache, - video_start, - audio_start, - ) - x0 = self.distribution_matching.x0_from_velocity( - block, - velocity, - sigma, - ) - if step_index < exit_index: - next_sigma = context.trajectory_scheduler.sigma_at( - step_index + 1, - device=self.model.device, - dtype=context.running_dtype, - ) - with torch.no_grad(): - noise = self.distribution_matching.random_noise_like( - x0, - x0.video.dtype, - lambda value: value, - ) - block = self.distribution_matching.add_noise( - context.trajectory_scheduler, - x0, - noise, - next_sigma, - ) - - video_chunks.append(x0.video) - audio_chunks.append(x0.audio) - cache_latents = self.distribution_matching.detach(x0) - cache_sigma = torch.full( - (1,), - context.context_noise, - device=self.model.device, - dtype=context.running_dtype, - ) - if context.context_noise > 0: - cache_latents = self.distribution_matching.add_noise( - context.trajectory_scheduler, - cache_latents, - self.distribution_matching.random_noise_like( - cache_latents, - cache_latents.video.dtype, - lambda value: value, - ), - cache_sigma, - ) - with torch.no_grad(): - self._predict_chunk( - cache_latents, - cache_sigma, - condition, - positions, - cache, - video_start, - audio_start, - ) - - generated = LTXJointLatents( - torch.cat(video_chunks, dim=1).to(context.running_dtype), - torch.cat(audio_chunks, dim=1).to(context.running_dtype), - latent_shape, - ) - return generated, exit_indices - - def _predict_chunk( - self, - latents, - sigma, - condition, - positions, - cache, - video_start, - audio_start, - ): - video_context, audio_context, context_mask = condition - video_positions, audio_positions = positions - if latents.video.shape[0] != 1 or latents.audio.shape[0] != 1: - raise ValueError("LTX autoregressive DMD only supports physical batch size 1.") - sigma = sigma.to( - device=self.model.device, - dtype=self.model.running_dtype, - ) - video = Modality( - enabled=True, - latent=latents.video.to( - device=self.model.device, - dtype=self.model.running_dtype, - ), - sigma=sigma, - timesteps=sigma.view(-1, 1) - .expand( - 1, - latents.video.shape[1], - ) - .clone(), - positions=video_positions.to(device=self.model.device), - context=video_context.to( - device=self.model.device, - dtype=self.model.running_dtype, - ), - context_mask=context_mask.to(device=self.model.device), - ) - audio = Modality( - enabled=True, - latent=latents.audio.to( - device=self.model.device, - dtype=self.model.running_dtype, - ), - sigma=sigma, - timesteps=sigma.view(-1, 1) - .expand( - 1, - latents.audio.shape[1], - ) - .clone(), - positions=audio_positions.to(device=self.model.device), - context=audio_context.to( - device=self.model.device, - dtype=self.model.running_dtype, - ), - context_mask=context_mask.to(device=self.model.device), - ) - with self.model.transformer_forward_context(): - video_prediction, audio_prediction = self.model.denoiser_module()( - video=video, - audio=audio, - perturbations=None, - kv_cache=cache, - video_current_start=video_start, - audio_current_start=audio_start, - ) - return LTXJointLatents( - video_prediction, - audio_prediction, - latents.shape, - ) - - def _new_caches( - self, - dtype, - device, - video_total_tokens, - audio_total_tokens, - video_frame_sequence_length, - num_frames, - frames_per_chunk, - ): - transformer = self._unwrap_transformer() - num_layers = len(transformer.transformer_blocks) - video_heads = int(transformer.num_attention_heads) - video_head_dim = int(transformer.inner_dim / video_heads) - audio_heads = int(transformer.audio_num_attention_heads) - audio_head_dim = int(transformer.audio_inner_dim / audio_heads) - local_attention_size = int(getattr(transformer, "local_attn_size", -1)) - sink_size = int(getattr(transformer, "sink_size", 0)) - video_chunk_tokens = frames_per_chunk * video_frame_sequence_length - if local_attention_size == -1: - video_cache_size = video_total_tokens - video_window = video_total_tokens - else: - video_window = max( - video_chunk_tokens, - local_attention_size * video_frame_sequence_length, - ) - video_cache_size = max( - video_window, - sink_size * video_frame_sequence_length + video_chunk_tokens, - ) - audio_tokens_per_frame = max( - 1, - math.ceil(audio_total_tokens / max(1, num_frames)), - ) - audio_chunk_tokens = max( - 1, - math.ceil(audio_total_tokens / (num_frames // frames_per_chunk)), - ) - if local_attention_size == -1: - audio_cache_size = audio_total_tokens - audio_window = audio_total_tokens - else: - audio_window = max( - audio_chunk_tokens, - local_attention_size * audio_tokens_per_frame, - ) - audio_cache_size = max( - audio_window, - sink_size * audio_tokens_per_frame + audio_chunk_tokens, - ) - return { - "video": [ - self._empty_cache( - video_cache_size, - video_heads, - video_head_dim, - dtype, - device, - video_window, - sink_size * video_frame_sequence_length, - ) - for _ in range(num_layers) - ], - "audio": [ - self._empty_cache( - audio_cache_size, - audio_heads, - audio_head_dim, - dtype, - device, - audio_window, - sink_size * audio_tokens_per_frame, - ) - for _ in range(num_layers) - ], - } - - @staticmethod - def _empty_cache( - cache_size, - heads, - head_dim, - dtype, - device, - attention_window_size, - sink_tokens, - ): - return { - "k": torch.zeros( - (1, cache_size, heads, head_dim), - dtype=dtype, - device=device, - ), - "v": torch.zeros( - (1, cache_size, heads, head_dim), - dtype=dtype, - device=device, - ), - "global_end_index": torch.zeros( - 1, - dtype=torch.long, - device=device, - ), - "local_end_index": torch.zeros( - 1, - dtype=torch.long, - device=device, - ), - "attention_window_size": int(attention_window_size), - "sink_tokens": int(sink_tokens), - } - - @staticmethod - def _audio_range(block_index, num_blocks, total_tokens): - start = (block_index * total_tokens + num_blocks - 1) // num_blocks - end = ((block_index + 1) * total_tokens + num_blocks - 1) // num_blocks - return start, end - - @staticmethod - def _sample_exit_indices(count, num_steps, device): - if torch.distributed.is_available() and torch.distributed.is_initialized(): - if torch.distributed.get_rank() == 0: - indices = torch.randint(0, num_steps, (count,), device=device) - else: - indices = torch.empty(count, dtype=torch.long, device=device) - torch.distributed.broadcast(indices, src=0) - return indices.tolist() - return torch.randint(0, num_steps, (count,), device=device).tolist() - - def _unwrap_transformer(self): - transformer = self.model.denoiser_module() - for candidate in ( - transformer, - getattr(transformer, "module", None), - getattr(transformer, "_fsdp_wrapped_module", None), - ): - if candidate is not None and hasattr(candidate, "_forward_inference"): - return candidate - return None diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py deleted file mode 100644 index caea3b7ce..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_consistency_model_capability.py +++ /dev/null @@ -1,194 +0,0 @@ -"""Consistency training for LTX's joint video/audio latent state.""" - -from __future__ import annotations - -import copy - -import torch -from torch import Tensor, nn - -from lightx2v_train.model_zoo.capability_adapters.consistency_model import ( - TimeConditionedConsistencyModelCapability, - TimeEmbeddingAdapter, -) -from lightx2v_train.model_zoo.native.ltx2 import Modality - -from .common import LTXLatentCodec - - -class LTXTimeEmbeddingAdapter(TimeEmbeddingAdapter): - """Condition both LTX modalities on a shared MeanFlow endpoint.""" - - def hook_modules(self, denoiser: nn.Module) -> tuple[nn.Module, ...]: - return denoiser.adaln_single, denoiser.audio_adaln_single - - def base_embedder(self, denoiser: nn.Module) -> nn.Module: - return denoiser.adaln_single - - def build_endpoint_embedder(self, denoiser: nn.Module) -> nn.Module: - return nn.ModuleDict( - { - "video": copy.deepcopy(denoiser.adaln_single), - "audio": copy.deepcopy(denoiser.audio_adaln_single), - } - ) - - def embedding_dimension(self, denoiser: nn.Module) -> int: - return int(denoiser.inner_dim) - - def encode( - self, - denoiser: nn.Module, - embedder: nn.Module, - time: Tensor, - ) -> Tensor: - parameter = next(embedder.parameters()) - scaled_time = time.to(device=parameter.device) * float(denoiser.timestep_scale_multiplier) - _, embedding = embedder( - scaled_time.flatten(), - hidden_dtype=parameter.dtype, - ) - return embedding - - def add_endpoint_embedding( - self, - denoiser: nn.Module, - endpoint_embedder: nn.ModuleDict, - endpoint_time: Tensor, - hook_module: nn.Module, - output, - ): - if not isinstance(output, tuple) or len(output) != 2 or not all(torch.is_tensor(value) for value in output): - raise TypeError("LTX AdaLN endpoint conditioning expects a pair of tensor outputs.") - if hook_module is denoiser.adaln_single: - branch = endpoint_embedder["video"] - elif hook_module is denoiser.audio_adaln_single: - branch = endpoint_embedder["audio"] - else: - raise RuntimeError("Received an endpoint hook from an unknown LTX AdaLN module.") - - parameter = next(branch.parameters()) - scaled_time = endpoint_time.to(device=parameter.device) * float(denoiser.timestep_scale_multiplier) - endpoint_outputs = branch( - scaled_time.flatten(), - hidden_dtype=parameter.dtype, - ) - return tuple(self._add_broadcast(value, endpoint) for value, endpoint in zip(output, endpoint_outputs)) - - @staticmethod - def _add_broadcast(value: Tensor, endpoint: Tensor) -> Tensor: - while endpoint.ndim < value.ndim: - endpoint = endpoint.unsqueeze(-2) - return value + endpoint.to(device=value.device, dtype=value.dtype) - - -class LTXConsistencyModelCapability(TimeConditionedConsistencyModelCapability): - """Expose joint LTX audio/video denoising through the tensor CM boundary.""" - - def __init__(self, model) -> None: - super().__init__(model, LTXTimeEmbeddingAdapter()) - config = model.config["training"].get("ltx2", {}) - self.latent_codec = LTXLatentCodec( - model, - default_fps=float(config.get("default_fps", 24.0)), - ) - - def encode_latent(self, batch): - _, _, _, _, video_tokens, audio_tokens = self.latent_codec.prepare( - batch["inputs"], - self.model.running_dtype, - ) - return self.latent_codec.flatten(video_tokens, audio_tokens) - - def encode_condition(self, batch): - conditioning = batch["conditioning"] - positive = conditioning.get("positive") - if positive is None: - video_context, audio_context, context_mask = self.model.encode_prompt_condition(conditioning.get("prompt", "")) - else: - video_context, audio_context, context_mask = self.model.prepare_text_condition(positive) - - video_data, _, video_latents, _, video_tokens, audio_tokens = self.latent_codec.prepare( - batch["inputs"], - self.model.running_dtype, - ) - video_positions, audio_positions = self.latent_codec.positions( - video_data, - video_latents, - audio_tokens, - ) - return { - "video_context": video_context, - "audio_context": audio_context, - "context_mask": context_mask, - "video_positions": video_positions, - "audio_positions": audio_positions, - "video_shape": tuple(video_tokens.shape), - "audio_shape": tuple(audio_tokens.shape), - } - - def sampling_latent_hw(self, batch, clean) -> tuple[int, int]: - del clean - _, _, video_latents, _, _, _ = self.latent_codec.prepare( - batch["inputs"], - self.model.running_dtype, - ) - return int(video_latents.shape[-2]), int(video_latents.shape[-1]) - - def predict(self, request, path): - with self._configured_prediction(request) as prepared_request: - if prepared_request.model_kwargs: - names = ", ".join(sorted(prepared_request.model_kwargs)) - raise TypeError(f"Unsupported LTX consistency model arguments: {names}.") - prediction = self._predict_velocity( - prepared_request.sample, - prepared_request.time, - prepared_request.condition, - ) - return path.convert_prediction( - prepared_request.sample, - prediction, - prepared_request.time, - source_type=self.model.denoiser_prediction_type(), - target_type=prepared_request.prediction_type, - ) - - def _predict_velocity(self, sample, time, condition): - video_tokens, audio_tokens = self.latent_codec.unflatten( - sample, - condition["video_shape"], - condition["audio_shape"], - ) - sigma = time.to(device=sample.device, dtype=sample.dtype).reshape(-1) - if sigma.numel() != 1: - raise ValueError(f"LTX consistency training requires one shared time value, got {tuple(time.shape)}.") - video_timesteps = sigma.view(1, 1).expand(1, video_tokens.shape[1]) - audio_timesteps = sigma.view(1, 1).expand(1, audio_tokens.shape[1]) - - video = Modality( - enabled=True, - latent=video_tokens, - sigma=sigma, - timesteps=video_timesteps, - positions=condition["video_positions"], - context=condition["video_context"], - context_mask=condition["context_mask"], - ) - audio = Modality( - enabled=True, - latent=audio_tokens, - sigma=sigma, - timesteps=audio_timesteps, - positions=condition["audio_positions"], - context=condition["audio_context"], - context_mask=condition["context_mask"], - ) - with self.model.transformer_forward_context(): - video_prediction, audio_prediction = self.denoiser()( - video=video, - audio=audio, - perturbations=None, - ) - video_prediction = video_prediction[:, -video_tokens.shape[1] :] - audio_prediction = audio_prediction[:, -audio_tokens.shape[1] :] - return self.latent_codec.flatten(video_prediction, audio_prediction) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py deleted file mode 100644 index 6c7df2a7a..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_distribution_matching_capability.py +++ /dev/null @@ -1,438 +0,0 @@ -"""Distribution-matching capability for LTX joint video/audio models.""" - -from __future__ import annotations - -import torch - -from lightx2v_train.model_zoo.capability_adapters.common import ( - GenericDistributionMatchingCapability, - _negative_prompt, - _require_single_prompt, - _require_singleton_tensor, -) -from lightx2v_train.model_zoo.native.ltx2 import ( - AudioLatentShape, - AudioPatchifier, - Modality, - SpatioTemporalScaleFactors, - VideoLatentPatchifier, - VideoLatentShape, - get_pixel_coords, -) -from lightx2v_train.utils.constants import LTX2_NEGATIVE_PROMPT - -from .common import LTXJointLatentShape, LTXJointLatents - - -class LTXDistributionMatchingCapability(GenericDistributionMatchingCapability): - """Distribution-matching operations for joint LTX video/audio states.""" - - def __init__(self, model) -> None: - super().__init__(model) - config = model.config["training"].get("ltx2", {}) - dmd_config = model.config["training"].get("dmd", {}) - self.video_weight = float(config.get("video_loss_weight", 1.0)) - self.audio_weight = float(config.get("audio_loss_weight", 1.0)) - self.default_fps = float(config.get("default_fps", dmd_config.get("fps", 25.0))) - self.video_patchifier = VideoLatentPatchifier(patch_size=1) - self.audio_patchifier = AudioPatchifier(patch_size=1) - - @property - def default_negative_prompt(self): - return LTX2_NEGATIVE_PROMPT - - @property - def default_lora_target_modules(self): - return None - - def latent_shape( - self, - batch, - shape_config, - image_sizes, - broadcast, - ): - del image_sizes, broadcast - prompt = batch["conditioning"].get("prompt", "") - _require_single_prompt(prompt) - video_shape = shape_config.get("video_latent_shape") - if video_shape is None: - meta = batch["meta"] - height = int(meta["target_height"][0].item()) if "target_height" in meta else int(shape_config.get("height", 512)) - width = int(meta["target_width"][0].item()) if "target_width" in meta else int(shape_config.get("width", 768)) - frames = int(shape_config.get("num_frames", 241)) - video_shape = ( - 1, - int(shape_config.get("video_latent_channels", 128)), - (frames - 1) // int(shape_config.get("video_temporal_scale", 8)) + 1, - height // int(shape_config.get("video_spatial_scale", 32)), - width // int(shape_config.get("video_spatial_scale", 32)), - ) - else: - video_shape = self._singleton_shape( - video_shape, - "video_latent_shape", - ) - - audio_shape = shape_config.get("audio_latent_shape") - if audio_shape is None: - if "audio_num_frames" in shape_config: - audio_frames = int(shape_config["audio_num_frames"]) - else: - pixel_frames = int( - shape_config.get( - "num_frames", - (video_shape[2] - 1) * 8 + 1, - ) - ) - fps = float(shape_config.get("fps", self.default_fps)) - audio_frames = round(pixel_frames / fps * 25.0) - audio_shape = ( - 1, - int(shape_config.get("audio_latent_channels", 8)), - audio_frames, - int(shape_config.get("audio_mel_bins", 16)), - ) - else: - audio_shape = self._singleton_shape( - audio_shape, - "audio_latent_shape", - ) - return LTXJointLatentShape( - video_latent=video_shape, - audio_latent=audio_shape, - video_tokens=( - 1, - video_shape[2] * video_shape[3] * video_shape[4], - video_shape[1], - ), - audio_tokens=( - 1, - audio_shape[2], - audio_shape[1] * audio_shape[3], - ), - fps=float(shape_config.get("fps", self.default_fps)), - ) - - def encode_conditions( - self, - batch, - negative_prompt, - guidance_scale, - broadcast, - ): - conditioning = batch["conditioning"] - positive_value = conditioning.get("positive") - if positive_value is None: - return super().encode_conditions( - batch, - negative_prompt, - guidance_scale, - broadcast, - ) - with torch.no_grad(): - positive = self._cached_condition(positive_value) - if guidance_scale > 1: - negative_value = conditioning.get("negative") - if negative_value is not None: - negative = self._cached_condition(negative_value) - else: - prompt = conditioning.get("prompt", "") - scalar = _require_single_prompt(prompt) - 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 _cached_condition(self, value): - if isinstance(value, (list, tuple)) and len(value) == 3: - video, audio, mask = value - elif isinstance(value, dict) and { - "video_context", - "audio_context", - "context_mask", - }.issubset(value): - video = value["video_context"] - audio = value["audio_context"] - mask = value["context_mask"] - else: - return self.model.prepare_text_condition(value) - video = _require_singleton_tensor( - video.to(device=self.device, dtype=self.model.running_dtype), - "LTX video context", - ) - audio = _require_singleton_tensor( - audio.to(device=self.device, dtype=self.model.running_dtype), - "LTX audio context", - ) - mask = _require_singleton_tensor( - mask.to(device=self.device), - "LTX context mask", - ) - return video, audio, mask - - def predict_velocity(self, latents, sigma, condition): - video_context, audio_context, context_mask = condition - _require_singleton_tensor(latents.video, "LTX video latent") - _require_singleton_tensor(latents.audio, "LTX audio latent") - sigma = sigma.to(device=self.device, dtype=self.model.running_dtype) - video_timesteps = ( - sigma.view(-1, 1) - .expand( - 1, - latents.video.shape[1], - ) - .clone() - ) - audio_timesteps = ( - sigma.view(-1, 1) - .expand( - 1, - latents.audio.shape[1], - ) - .clone() - ) - video = Modality( - enabled=True, - latent=latents.video.to( - device=self.device, - dtype=self.model.running_dtype, - ), - sigma=sigma, - timesteps=video_timesteps, - positions=self._video_token_positions(latents.shape), - context=video_context.to( - device=self.device, - dtype=self.model.running_dtype, - ), - context_mask=context_mask.to(device=self.device), - ) - audio = Modality( - enabled=True, - latent=latents.audio.to( - device=self.device, - dtype=self.model.running_dtype, - ), - sigma=sigma, - timesteps=audio_timesteps, - positions=self._audio_token_positions(latents.shape), - context=audio_context.to( - device=self.device, - dtype=self.model.running_dtype, - ), - context_mask=context_mask.to(device=self.device), - ) - with self.model.transformer_forward_context(): - video_prediction, audio_prediction = self.model.denoiser_module()( - video=video, - audio=audio, - perturbations=None, - ) - return LTXJointLatents( - video_prediction, - audio_prediction, - latents.shape, - ) - - def predict_guided_velocity( - self, - latents, - sigma, - condition, - negative_condition, - guidance_scale, - cfg_norm, - ): - positive = self.predict_velocity(latents, sigma, condition) - if negative_condition is None: - return positive - negative = self.predict_velocity(latents, sigma, negative_condition) - return LTXJointLatents( - self._cfg( - positive.video, - negative.video, - guidance_scale, - cfg_norm, - ), - self._cfg( - positive.audio, - negative.audio, - guidance_scale, - cfg_norm, - ), - latents.shape, - ) - - def initial_latents(self, latent_shape, dtype, broadcast): - del broadcast - if latent_shape.video_tokens[0] != 1 or latent_shape.audio_tokens[0] != 1: - raise ValueError(f"LTX DMD latent shape must start with 1, got {latent_shape}.") - return LTXJointLatents( - torch.randn( - latent_shape.video_tokens, - device=self.device, - dtype=dtype, - ), - torch.randn( - latent_shape.audio_tokens, - device=self.device, - dtype=dtype, - ), - latent_shape, - ) - - @staticmethod - def latent_hw(latent_shape): - return None - - @staticmethod - def random_noise_like(latents, dtype, broadcast): - del broadcast - return LTXJointLatents( - torch.randn_like(latents.video, dtype=dtype), - torch.randn_like(latents.audio, dtype=dtype), - latents.shape, - ) - - @staticmethod - def add_noise(scheduler, latents, noise, sigma): - return LTXJointLatents( - scheduler.add_noise(latents.video, noise.video, sigma).to(latents.video.dtype), - scheduler.add_noise(latents.audio, noise.audio, sigma).to(latents.audio.dtype), - latents.shape, - ) - - @staticmethod - def training_target(latents, noise): - return LTXJointLatents( - noise.video - latents.video.float(), - noise.audio - latents.audio.float(), - latents.shape, - ) - - @staticmethod - def step(scheduler, velocity, step_index, sample): - video_next, video_x0 = scheduler.step_by_index( - velocity.video, - step_index, - sample.video, - ) - audio_next, audio_x0 = scheduler.step_by_index( - velocity.audio, - step_index, - sample.audio, - ) - return ( - LTXJointLatents(video_next, audio_next, sample.shape), - LTXJointLatents(video_x0, audio_x0, sample.shape), - ) - - @staticmethod - def x0_from_velocity(sample, velocity, sigma): - def project(value, flow): - expanded = sigma.reshape( - sigma.shape[0], - *([1] * (value.ndim - 1)), - ) - return value - expanded * flow - - return LTXJointLatents( - project(sample.video, velocity.video), - project(sample.audio, velocity.audio), - sample.shape, - ) - - def regression_loss(self, prediction, target): - video = torch.nn.functional.mse_loss( - prediction.video.float(), - target.video.float(), - ) - audio = torch.nn.functional.mse_loss( - prediction.audio.float(), - target.audio.float(), - ) - return self.video_weight * video + self.audio_weight * audio - - 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 LTXJointLatents( - value.video.detach(), - value.audio.detach(), - value.shape, - ) - - @staticmethod - def to_dtype(value, dtype): - return LTXJointLatents( - value.video.to(dtype=dtype), - value.audio.to(dtype=dtype), - value.shape, - ) - - def extract_real_latents(self, batch, dtype, broadcast): - raise ValueError("LTX joint DMD does not support real-data fake loss.") - - def _video_token_positions(self, shape): - _, channels, frames, height, width = shape.video_latent - latent_coords = self.video_patchifier.get_patch_grid_bounds( - output_shape=VideoLatentShape( - frames=frames, - height=height, - width=width, - batch=shape.video_latent[0], - channels=channels, - ), - device=self.device, - ) - positions = get_pixel_coords( - latent_coords=latent_coords, - scale_factors=SpatioTemporalScaleFactors.default(), - causal_fix=True, - ).float() - positions[:, 0, ...] /= shape.fps - return positions - - def _audio_token_positions(self, shape): - _, channels, frames, mel_bins = shape.audio_latent - return self.audio_patchifier.get_patch_grid_bounds( - output_shape=AudioLatentShape( - frames=frames, - mel_bins=mel_bins, - batch=shape.audio_latent[0], - channels=channels, - ), - device=self.device, - ) - - @staticmethod - def _singleton_shape(shape, name): - shape = tuple(int(dimension) for dimension in shape) - expected = 4 if name.startswith("audio") else 5 - without_batch = expected - 1 - if len(shape) == without_batch: - return (1, *shape) - if len(shape) == expected: - if shape[0] != 1: - raise ValueError(f"training.dmd.{name} must begin with singleton dimension 1, got {shape}.") - return shape - raise ValueError(f"training.dmd.{name} has unsupported shape {shape}.") diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py deleted file mode 100644 index b1d684a6e..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_flow_matching_capability.py +++ /dev/null @@ -1,143 +0,0 @@ -"""Flow-matching capability for LTX joint video/audio models.""" - -from __future__ import annotations - -from typing import Any, Mapping - -import torch -from torch import Tensor - -from lightx2v_train.model_capabilities import ( - BoundCapability, - FlowMatchingSFTCapability, - LossResult, - SFTStepContext, -) -from lightx2v_train.model_zoo.native.ltx2 import Modality - -from .common import LTXLatentCodec - - -class LTXFlowMatchingCapability(BoundCapability, FlowMatchingSFTCapability): - """Supervised flow matching for LTX's joint video/audio denoiser.""" - - def __init__(self, model) -> None: - super().__init__(model) - config = model.config["training"].get("ltx2", {}) - self.video_loss_weight = float(config.get("video_loss_weight", 1.0)) - self.audio_loss_weight = float(config.get("audio_loss_weight", 1.0)) - self.default_fps = float(config.get("default_fps", 24.0)) - self.latent_codec = LTXLatentCodec(model, default_fps=self.default_fps) - - def compute_loss( - self, - batch: Mapping[str, Any], - context: SFTStepContext, - ) -> LossResult: - inputs = batch["inputs"] - conditions = batch["conditioning"].get("positive") - if conditions is None: - raise KeyError("LTX flow matching expects conditioning.positive.") - video_data, _, video_latents, _, video_tokens, audio_tokens = self.latent_codec.prepare( - inputs, - context.running_dtype, - ) - - video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) - video_noisy = self._initialize_noisy_target( - video_tokens, - context.noise_scheduler, - ) - audio_noisy = self._initialize_noisy_target( - audio_tokens, - context.noise_scheduler, - sigmas=video_noisy[-1], - ) - video_input, video_target, video_timesteps, video_mask, video_sigmas = video_noisy - audio_input, audio_target, audio_timesteps, audio_mask, audio_sigmas = audio_noisy - - video_positions, audio_positions = self.latent_codec.positions( - video_data, - video_latents, - audio_tokens, - ) - video = Modality( - enabled=True, - latent=video_input, - sigma=video_sigmas, - timesteps=video_timesteps, - positions=video_positions, - context=video_context, - context_mask=context_mask, - ) - audio = Modality( - enabled=True, - latent=audio_input, - sigma=audio_sigmas, - timesteps=audio_timesteps, - positions=audio_positions, - context=audio_context, - context_mask=context_mask, - ) - - with self.model.transformer_forward_context(): - video_prediction, audio_prediction = self.model.denoiser_module()( - video=video, - audio=audio, - perturbations=None, - ) - - video_loss = self._masked_mse( - video_prediction, - video_target, - video_mask, - ).mean() - audio_loss = self._masked_mse( - audio_prediction, - audio_target, - audio_mask, - ).mean() - loss = self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss - return LossResult( - loss=loss, - metrics={ - "video_loss": video_loss.detach(), - "audio_loss": audio_loss.detach(), - }, - ) - - @staticmethod - def _initialize_noisy_target(latents: Tensor, scheduler, sigmas=None): - _, sequence_length, _ = latents.shape - if sigmas is None: - sigmas = scheduler.sample_timestep_or_sigma( - seq_len=sequence_length, - ) - sigmas = sigmas.to(device=latents.device, dtype=latents.dtype) - noise = torch.randn_like(latents) - expanded = sigmas.view(-1, 1, 1) - noisy = (1 - expanded) * latents + expanded * noise - target = noise - latents - timesteps = ( - sigmas.view(-1, 1) - .expand( - 1, - sequence_length, - ) - .clone() - ) - mask = torch.ones( - 1, - sequence_length, - dtype=torch.bool, - device=latents.device, - ) - return noisy, target, timesteps, mask, sigmas - - @staticmethod - def _masked_mse(prediction, target, mask): - target_length = target.shape[1] - prediction = prediction[:, -target_length:, :] - weights = mask[:, -target_length:].unsqueeze(-1).float() - loss = (prediction.float() - target.float()).square() * weights - return loss.mean(dim=(-2, -1)) / weights.mean(dim=(-2, -1)).clamp_min(1e-8) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py deleted file mode 100644 index f1c7767b7..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/capability_adapters/ltx_teacher_forcing_capability.py +++ /dev/null @@ -1,357 +0,0 @@ -"""Teacher-forcing capability for LTX models.""" - -from __future__ import annotations - -from typing import Any, Mapping - -import torch - -from lightx2v_train.model_capabilities import ( - LossResult, - TeacherForcingCapability, - TeacherForcingStepContext, -) -from lightx2v_train.model_zoo.native.ltx2 import AudioLatentShape, Modality -from lightx2v_train.runtime.distributed import is_sequence_parallel_enabled - -from .ltx_flow_matching_capability import LTXFlowMatchingCapability - - -class LTXTeacherForcingCapability( - LTXFlowMatchingCapability, - TeacherForcingCapability, -): - """Chunk-wise teacher forcing for LTX's joint video/audio denoiser.""" - - def __init__(self, model) -> None: - super().__init__(model) - self._mask_cache = {} - - def compute_loss( - self, - batch: Mapping[str, Any], - context: TeacherForcingStepContext, - ) -> LossResult: - if is_sequence_parallel_enabled(): - raise ValueError("LTX teacher forcing does not support sequence parallel.") - - inputs = batch["inputs"] - video_data = inputs.get("video_latents") - audio_data = inputs.get("audio_latents") - conditions = batch["conditioning"].get("positive") - if video_data is None or audio_data is None or conditions is None: - raise KeyError("LTX teacher forcing expects inputs.video_latents, inputs.audio_latents, and conditioning.positive.") - - with torch.no_grad(): - video_latents = self._video_latents( - video_data, - context.running_dtype, - ) - audio_latents = self._audio_latents( - audio_data, - context.running_dtype, - ) - video_latents, audio_latents = self._crop_to_chunk_multiple( - video_latents, - audio_latents, - context.num_frame_per_chunk, - ) - num_video_frames = int(video_latents.shape[2]) - video_tokens = self.video_patchifier.patchify(video_latents) - audio_tokens = self.audio_patchifier.patchify(audio_latents) - video_context, audio_context, context_mask = self.model.prepare_text_condition(conditions) - - video_inputs = self._prepare_video_inputs( - video_tokens, - num_frames=num_video_frames, - spatial_tokens=int(video_latents.shape[3] * video_latents.shape[4]), - context=context, - ) - ( - video_input, - video_target, - video_timesteps, - video_weights, - video_sigmas, - video_frame_weights, - video_augmentation_sigmas, - video_attention_mask, - video_block_ids, - ) = video_inputs - audio_inputs = self._prepare_audio_inputs( - audio_tokens, - frame_sigmas=video_sigmas, - frame_weights=video_frame_weights, - augmentation_frame_sigmas=video_augmentation_sigmas, - num_video_frames=num_video_frames, - frames_per_chunk=context.num_frame_per_chunk, - ) - ( - audio_input, - audio_target, - audio_timesteps, - audio_weights, - audio_sigmas, - audio_attention_mask, - audio_block_ids, - ) = audio_inputs - - video_positions = self._video_positions( - video_data, - video_latents, - ) - video_positions = torch.cat( - [video_positions, video_positions], - dim=2, - ) - audio_positions = self.audio_patchifier.get_patch_grid_bounds( - output_shape=AudioLatentShape( - frames=audio_tokens.shape[1], - mel_bins=16, - batch=1, - channels=8, - ), - device=audio_tokens.device, - ) - audio_positions = torch.cat( - [audio_positions, audio_positions], - dim=2, - ) - video_cross_mask = self._cross_attention_mask( - video_tokens.device, - video_block_ids, - audio_block_ids, - ) - audio_cross_mask = self._cross_attention_mask( - audio_tokens.device, - audio_block_ids, - video_block_ids, - ) - - video = Modality( - enabled=True, - latent=video_input, - sigma=video_sigmas[:, 0], - timesteps=video_timesteps, - positions=video_positions, - context=video_context, - context_mask=context_mask, - attention_mask=video_attention_mask, - cross_attention_mask=video_cross_mask, - ) - audio = Modality( - enabled=True, - latent=audio_input, - sigma=audio_sigmas[:, 0], - timesteps=audio_timesteps, - positions=audio_positions, - context=audio_context, - context_mask=context_mask, - attention_mask=audio_attention_mask, - cross_attention_mask=audio_cross_mask, - ) - with self.model.transformer_forward_context(): - video_prediction, audio_prediction = self.model.denoiser_module()( - video=video, - audio=audio, - perturbations=None, - ) - - video_loss = self._weighted_token_mse( - video_prediction, - video_target, - video_weights, - ).mean() - audio_loss = self._weighted_token_mse( - audio_prediction, - audio_target, - audio_weights, - ).mean() - loss = self.video_loss_weight * video_loss + self.audio_loss_weight * audio_loss - return LossResult( - loss=loss, - metrics={ - "video_loss": video_loss.detach(), - "audio_loss": audio_loss.detach(), - }, - ) - - @staticmethod - def _crop_to_chunk_multiple(video, audio, frames_per_chunk): - if frames_per_chunk <= 1: - return video, audio - num_frames = int(video.shape[2]) - keep_frames = num_frames // frames_per_chunk * frames_per_chunk - if keep_frames == num_frames: - return video, audio - if keep_frames <= 0: - raise ValueError(f"LTX sample has {num_frames} latent frames, fewer than frames_per_chunk={frames_per_chunk}.") - audio_frames = int(audio.shape[2]) - keep_audio_frames = max(1, audio_frames * keep_frames // num_frames) - return ( - video[:, :, :keep_frames].contiguous(), - audio[:, :, :keep_audio_frames].contiguous(), - ) - - def _prepare_video_inputs( - self, - tokens, - num_frames, - spatial_tokens, - context, - ): - frame_sigmas, frame_weights = context.scheduler.sample_chunkwise( - num_frames=num_frames, - num_frame_per_chunk=context.num_frame_per_chunk, - device=tokens.device, - dtype=tokens.dtype, - ) - token_sigmas = frame_sigmas.repeat_interleave(spatial_tokens, dim=1) - token_weights = frame_weights.repeat_interleave(spatial_tokens, dim=1) - noisy, target, noise = self._add_token_noise(tokens, token_sigmas) - clean = tokens - clean_timesteps = torch.zeros_like(token_sigmas) - augmentation_frame_sigmas = None - if context.noise_augmentation_max_timestep > 0: - augmentation_frame_sigmas = context.scheduler.sample_clean_augmentation( - num_frames=num_frames, - num_frame_per_chunk=context.num_frame_per_chunk, - max_timestep=context.noise_augmentation_max_timestep, - device=tokens.device, - dtype=tokens.dtype, - ) - clean_timesteps = augmentation_frame_sigmas.repeat_interleave( - spatial_tokens, - dim=1, - ) - clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise - frame_ids = torch.arange( - num_frames, - device=tokens.device, - ).repeat_interleave(spatial_tokens) - block_ids = frame_ids // context.num_frame_per_chunk - attention_mask = self._attention_mask(tokens.device, block_ids) - return ( - torch.cat([clean, noisy], dim=1), - target, - torch.cat([clean_timesteps, token_sigmas], dim=1), - token_weights, - frame_sigmas, - frame_weights, - augmentation_frame_sigmas, - attention_mask, - block_ids, - ) - - def _prepare_audio_inputs( - self, - tokens, - frame_sigmas, - frame_weights, - augmentation_frame_sigmas, - num_video_frames, - frames_per_chunk, - ): - _, sequence_length, _ = tokens.shape - frame_ids = (torch.arange(sequence_length, device=tokens.device) * max(1, num_video_frames) // max(1, sequence_length)).clamp(max=frame_sigmas.shape[1] - 1) - gather_ids = frame_ids.unsqueeze(0) - token_sigmas = frame_sigmas.gather(1, gather_ids) - token_weights = frame_weights.gather(1, gather_ids) - noisy, target, noise = self._add_token_noise(tokens, token_sigmas) - clean = tokens - clean_timesteps = torch.zeros_like(token_sigmas) - if augmentation_frame_sigmas is not None: - clean_timesteps = augmentation_frame_sigmas.gather(1, gather_ids) - clean = (1.0 - clean_timesteps.unsqueeze(-1)) * tokens + clean_timesteps.unsqueeze(-1) * noise - block_ids = frame_ids // frames_per_chunk - attention_mask = self._attention_mask(tokens.device, block_ids) - return ( - torch.cat([clean, noisy], dim=1), - target, - torch.cat([clean_timesteps, token_sigmas], dim=1), - token_weights, - token_sigmas, - attention_mask, - block_ids, - ) - - def _attention_mask(self, device, block_ids): - key = ("self", self._device_key(device), self._block_key(block_ids)) - if key not in self._mask_cache: - self._mask_cache[key] = self._call_mask_builder( - "prepare_teacher_forcing_mask_from_block_ids", - device, - block_ids, - ) - return self._mask_cache[key] - - def _cross_attention_mask(self, device, query_ids, key_ids): - key = ( - "cross", - self._device_key(device), - self._block_key(query_ids), - self._block_key(key_ids), - ) - if key not in self._mask_cache: - self._mask_cache[key] = self._call_mask_builder( - "prepare_teacher_forcing_cross_mask_from_block_ids", - device, - query_ids, - key_ids, - ) - return self._mask_cache[key] - - def _call_mask_builder(self, method_name, *args): - transformer = self.model.denoiser_module() - candidates = ( - transformer, - getattr(transformer, "module", None), - getattr(transformer, "_fsdp_wrapped_module", None), - ) - for candidate in candidates: - method = getattr(candidate, method_name, None) - if method is not None: - return method(*args) - raise RuntimeError(f"LTX transformer does not implement {method_name}().") - - @staticmethod - def _device_key(device): - device = torch.device(device) - index = device.index - if device.type == "cuda" and index is None and torch.cuda.is_available(): - index = torch.cuda.current_device() - return device.type, index - - @staticmethod - def _block_key(block_ids): - values = block_ids.detach().to(dtype=torch.long) - if values.numel() == 0: - return 0, None, None, None, None, None - positions = torch.arange( - values.numel(), - device=values.device, - dtype=torch.long, - ) - return ( - int(values.numel()), - int(values[0].item()), - int(values[-1].item()), - int(values.max().item()), - int(values.sum().item()), - int((values * positions).sum().item()), - ) - - @staticmethod - def _add_token_noise(tokens, sigmas, noise=None): - if noise is None: - noise = torch.randn_like(tokens) - expanded = sigmas.unsqueeze(-1) - noisy = (1.0 - expanded) * tokens + expanded * noise - return noisy, noise - tokens, noise - - @staticmethod - def _weighted_token_mse(prediction, target, weights): - prediction = prediction[:, -target.shape[1] :, :] - loss = (prediction.float() - target.float()).square().mean(dim=-1) - weights = weights.float() - return (loss * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1e-8) diff --git a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py b/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py deleted file mode 100644 index 343c430d1..000000000 --- a/lightx2v_train/lightx2v_train/model_zoo/ltx/ltx_t2av.py +++ /dev/null @@ -1,416 +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_capabilities import ( - AutoregressiveDistributionMatchingCapability, - ConsistencyModelCapability, - DistributionMatchingCapability, - FlowMatchingSFTCapability, - TeacherForcingCapability, -) -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_autoregressive_distribution_matching_capability import ( - LTXAutoregressiveDistributionMatchingCapability, -) -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_consistency_model_capability import LTXConsistencyModelCapability -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_distribution_matching_capability import LTXDistributionMatchingCapability -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_flow_matching_capability import LTXFlowMatchingCapability -from lightx2v_train.model_zoo.ltx.capability_adapters.ltx_teacher_forcing_capability import LTXTeacherForcingCapability -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 register_capabilities(self): - super().register_capabilities() - self.capabilities.register( - FlowMatchingSFTCapability, - LTXFlowMatchingCapability(self), - ) - self.capabilities.register( - DistributionMatchingCapability, - LTXDistributionMatchingCapability(self), - ) - self.capabilities.register( - ConsistencyModelCapability, - LTXConsistencyModelCapability(self), - ) - if self.config["model"]["name"] == "ltx_t2av_ar": - self.capabilities.register( - AutoregressiveDistributionMatchingCapability, - LTXAutoregressiveDistributionMatchingCapability(self), - ) - self.capabilities.register( - TeacherForcingCapability, - LTXTeacherForcingCapability(self), - ) - - def load_components( - self, - *, - load_transformer, - load_vae, - load_condition_encoder, - ): - del load_vae - 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")) - should_load_transformer = load_transformer and bool(model_config.get("load_transformer", True)) - should_load_embeddings_processor = load_condition_encoder and bool(model_config.get("load_embeddings_processor", True)) - should_load_text_encoder = load_condition_encoder and 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)) - - transformer_path = self._resolve_transformer_path(model_path, model_config) - if should_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 should_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 should_load_text_encoder: - self.embeddings_processor.feature_extractor = None - if should_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 reuse_frozen_components_from(self, source): - super().reuse_frozen_components_from(source) - self.embeddings_processor = source.embeddings_processor - self.text_encoder = getattr(source, "text_encoder", None) - self.text_encoder_path = getattr(source, "text_encoder_path", self.text_encoder_path) - self.text_encoder_cpu = getattr(source, "text_encoder_cpu", self.text_encoder_cpu) - - 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): - self.loaded_student_checkpoint_path = None - if 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) - if len(prompts) != 1: - raise ValueError(f"LTX2 training requires exactly one prompt, got {len(prompts)}.") - 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/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: