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..df35d1788 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/config.py @@ -0,0 +1,362 @@ +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..ea21cb6ed --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/consistency/pcm.py @@ -0,0 +1,226 @@ +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(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 + + 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 \