diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json index b275ad7df..5c52a73ec 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json index 4632e3cec..f9db2b856 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json index f78e192dd..ad427bfff 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": false, "vae_decode_parallel": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json index dea873de9..19819f2e2 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json index 0bd223eb2..035390763 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json @@ -25,6 +25,7 @@ "use_compile": true, "warmup": true, "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "vae_use_compile": true, "vae_attn_type": "sage_attn2", "video_flow_shift": 6.0, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json index c51f93b18..2a65cffc2 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "dit_prepost_resident": true, "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json index 5baccfcfd..995bd43eb 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json @@ -37,6 +37,7 @@ "use_compile": true, "warmup": true, "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "vae_use_compile": true, "vae_encode_parallel": true, "vae_decode_tile_shape": { diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json index 788dfd3c0..e1feb03aa 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": false, "vae_decode_parallel": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json index 9bd651067..a76bbe433 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": false, "vae_decode_parallel": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json index c79948621..c2e603c37 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": false, "vae_decode_parallel": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json index de146a14d..42e082753 100755 --- a/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json @@ -8,6 +8,8 @@ "cpu_offload": true, "reference_image_resize_mode": "match", "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, "lazy_load": false, diff --git a/configs/minimax_h3/fp8/minimax_h3.json b/configs/minimax_h3/fp8/minimax_h3.json index 0d743f2f1..61a955ed9 100644 --- a/configs/minimax_h3/fp8/minimax_h3.json +++ b/configs/minimax_h3/fp8/minimax_h3.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, "lazy_load": false, diff --git a/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json index 29febcf8c..b23bcd897 100644 --- a/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_quantized": true, "text_encoder_quant_scheme": "fp8-sgl", diff --git a/configs/minimax_h3/fp8/minimax_h3_sp_5090.json b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json index 5757f4bfb..82e9268fd 100644 --- a/configs/minimax_h3/fp8/minimax_h3_sp_5090.json +++ b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": false, diff --git a/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json index ebdaa25ea..c83671b84 100644 --- a/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/minimax_h3.json b/configs/minimax_h3/minimax_h3.json index 660248780..6d956e57c 100644 --- a/configs/minimax_h3/minimax_h3.json +++ b/configs/minimax_h3/minimax_h3.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, "lazy_load": false, diff --git a/configs/minimax_h3/minimax_h3_block_offload.json b/configs/minimax_h3/minimax_h3_block_offload.json index 8bd0ddcf2..73b84dfbb 100644 --- a/configs/minimax_h3/minimax_h3_block_offload.json +++ b/configs/minimax_h3/minimax_h3_block_offload.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/minimax_h3_compile.json b/configs/minimax_h3/minimax_h3_compile.json index 0ea702fcf..942467485 100644 --- a/configs/minimax_h3/minimax_h3_compile.json +++ b/configs/minimax_h3/minimax_h3_compile.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, "lazy_load": false, diff --git a/configs/minimax_h3/minimax_h3_sol_block_offload.json b/configs/minimax_h3/minimax_h3_sol_block_offload.json index 70d92eb47..8f60bc649 100644 --- a/configs/minimax_h3/minimax_h3_sol_block_offload.json +++ b/configs/minimax_h3/minimax_h3_sol_block_offload.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "vae_cpu_offload": true, diff --git a/configs/minimax_h3/minimax_h3_sp.json b/configs/minimax_h3/minimax_h3_sp.json index 71b57b82c..3bd107ecc 100644 --- a/configs/minimax_h3/minimax_h3_sp.json +++ b/configs/minimax_h3/minimax_h3_sp.json @@ -7,6 +7,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, "vae_decode_parallel": true, diff --git a/configs/minimax_h3/minimax_h3_tp.json b/configs/minimax_h3/minimax_h3_tp.json index 88f0e533f..d952b11a9 100644 --- a/configs/minimax_h3/minimax_h3_tp.json +++ b/configs/minimax_h3/minimax_h3_tp.json @@ -6,6 +6,8 @@ "fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "offload_granularity": "model", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, diff --git a/configs/minimax_h3/minimax_h3_tp_sp.json b/configs/minimax_h3/minimax_h3_tp_sp.json index c7c10025d..9eb0e7cd4 100644 --- a/configs/minimax_h3/minimax_h3_tp_sp.json +++ b/configs/minimax_h3/minimax_h3_tp_sp.json @@ -6,6 +6,8 @@ "fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "offload_granularity": "model", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json index bd8d77607..52b7a97d3 100644 --- a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json +++ b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_15s.json @@ -6,6 +6,8 @@ "fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "offload_granularity": "model", "text_encoder_cpu_offload": true, "text_encoder_tensor_parallel": false, diff --git a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json index 623d98623..7be7f1782 100644 --- a/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json +++ b/configs/platforms/ascend_npu/minimax_h3_t2av_sp_compile_5s.json @@ -6,6 +6,8 @@ "fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "offload_granularity": "model", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": false, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json index 4bf1738c9..20e2b9dea 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_dmd_lora_4step_sp_tp.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json index 42f1ce43d..3901e1b83 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_int8_sp_tp.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_tensor_parallel": true, "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json index 837a18206..06f7ee092 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_tensor_parallel": true, "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json index 541897853..5032d12e5 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_sp_tp_cpu_offload.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_tensor_parallel": true, diff --git a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json index ae4b580f1..fec9022fa 100644 --- a/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json +++ b/configs/platforms/intel_xpu/dist_infer/minimax_h3_t2av_tp.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_tensor_parallel": true, "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json index ea60dc55b..c402cbbd3 100644 --- a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av.json b/configs/platforms/intel_xpu/minimax_h3_t2av.json index 82c499443..4864833ec 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json b/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json index 72f53c852..784582a6a 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_dmd_lora_4step.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json b/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json index 910bc1090..752f4bcd7 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_fp8.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json b/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json index a9ad4848b..92da77646 100644 --- a/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json +++ b/configs/platforms/intel_xpu/minimax_h3_t2av_int8.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_host_pinned": false, diff --git a/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json b/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json index b723935ce..8a4c92606 100644 --- a/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json +++ b/configs/platforms/metax/minimax_h3_t2av_tp1_block_offload.json @@ -23,6 +23,7 @@ "use_compile": true, "warmup": true, "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "vae_use_compile": false, "video_flow_shift": 12.0, "audio_flow_shift": 3.0, diff --git a/configs/platforms/metax/minimax_h3_t2av_tp_sp.json b/configs/platforms/metax/minimax_h3_t2av_tp_sp.json index eac89a061..05b4edbe2 100644 --- a/configs/platforms/metax/minimax_h3_t2av_tp_sp.json +++ b/configs/platforms/metax/minimax_h3_t2av_tp_sp.json @@ -20,6 +20,7 @@ "use_compile": true, "warmup": true, "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "vae_use_compile": true, "video_flow_shift": 12.0, "audio_flow_shift": 3.0, diff --git a/configs/platforms/mlu/minimax_h3_t2av_sp.json b/configs/platforms/mlu/minimax_h3_t2av_sp.json index dfe81eeef..104527e08 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_sp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_sp.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, "vae_cpu_offload": false, diff --git a/configs/platforms/mlu/minimax_h3_t2av_tp.json b/configs/platforms/mlu/minimax_h3_t2av_tp.json index ea9dfa860..dd2cd0361 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_tp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_tp.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, "vae_cpu_offload": false, diff --git a/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json b/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json index f3341f721..5ecb64a84 100644 --- a/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json +++ b/configs/platforms/mlu/minimax_h3_t2av_tp_sp.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, "vae_cpu_offload": false, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av.json b/configs/platforms/mthreads_musa/minimax_h3_t2av.json index c40e57d2d..7232693ad 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av.json @@ -8,6 +8,8 @@ "enable_cfg": false, "cpu_offload": true, "offload_granularity": "block", + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", "text_encoder_tensor_parallel": false, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json index d5a0b2f33..8ad59551c 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_fp8.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": false, "text_encoder_quantized": true, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json index 9842643a2..e1845861d 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, "vae_cpu_offload": false, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json index a60275c76..d5e0e6a0f 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp4_sp2.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "offload_granularity": "model", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json index de92bce18..18d39f57b 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp_fp8.json @@ -7,6 +7,8 @@ "target_fps": 24, "enable_cfg": false, "cpu_offload": false, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": false, "text_encoder_tensor_parallel": true, "text_encoder_quantized": true, diff --git a/lightx2v/models/networks/minimax_h3/adaln_cache.py b/lightx2v/models/networks/minimax_h3/adaln_cache.py new file mode 100644 index 000000000..2bdce39a3 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/adaln_cache.py @@ -0,0 +1,199 @@ +"""Identify, validate, and load persistent MiniMax-H3 AdaLN caches. + +Offline generation lives in ``tools/cache_minimax_h3_adaln/builder.py`` so the +inference path does not carry checkpoint-building concerns. +""" + +import json +from pathlib import Path + +import torch +from loguru import logger +from safetensors import SafetensorError, safe_open + +from lightx2v.models.networks.minimax_h3.packing import ( + CONDITION_AUDIO_TIMESTEP, + KEYFRAME_NOISE_AUG, +) +from lightx2v.models.schedulers.minimax_h3.scheduler import _make_schedule + + +def validate_adaln_cache_config(config) -> None: + """Reject modes whose AdaLN result cannot be represented by this cache.""" + if not config.get("use_adaln_cache", False): + return + cache_dir = config.get("adaln_cache_dir") + if cache_dir is None or not str(cache_dir).strip(): + separator = "=" * 88 + message = ( + f"\n{separator}\n" + "MINIMAX-H3 ADALN CACHE CONFIGURATION ERROR\n" + "use_adaln_cache=true, but adaln_cache_dir is missing or empty.\n" + "\nACTION REQUIRED\n" + "Add the cache root to the inference JSON config:\n" + ' "use_adaln_cache": true,\n' + ' "adaln_cache_dir": "~/.cache/lightx2v/adaln"\n' + "adaln_cache_dir can be any custom cache root; the path above is the recommended default.\n" + "\nBefore inference, set lightx2v_path, model_path, --config_json, and " + "--task in tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh.\n" + "Then generate the cache from the repository root with:\n" + " bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh\n" + "Set --task fl2av in that script for t2av/i2av/l2av/fl2av, or " + "--task ref2av for ref2av. The generation script and inference must " + f"use the same JSON config.\n{separator}" + ) + logger.error(message) + raise ValueError(message) + if config.get("dummy_model", False): + raise NotImplementedError("Persistent MiniMax-H3 AdaLN cache does not support dummy_model") + + +def _cache_root(config) -> Path: + return Path(config["adaln_cache_dir"]).expanduser().resolve() + + +def _selected_profiles(config) -> list[str]: + task = str(config.get("task", "")) + if task == "ref2av": + # Ref2AV always has visual reference rows and may additionally have + # frozen audio rows. It uses transformer_ref, so its cache must remain + # separate from every base-transformer task. + return ["ref2av_video", "ref2av_video_audio"] + if task in {"t2av", "i2av", "l2av", "fl2av"}: + # All base-transformer tasks share this pair, so one FL2AV cache also + # serves T2AV, I2AV, and L2AV without support_tasks-dependent paths. + return ["t2av", "conditioned"] + raise ValueError(f"No persistent AdaLN cache profile is available for task: {task!r}") + + +def _float32_bits(values) -> list[int]: + # Keep timestep keys bit-exact across JSON serialization and later runs. + return torch.tensor(list(values), dtype=torch.float32).view(torch.int32).tolist() + + +def _timesteps_from_bits(bits: list[int], device="cpu") -> torch.Tensor: + return torch.tensor(bits, dtype=torch.int32).view(torch.float32).to(device) + + +def _cache_entries(config, profiles: list[str]) -> list[dict]: + infer_steps = int(config["infer_steps"]) + _, video_timesteps = _make_schedule( + infer_steps, + float(config.get("video_flow_shift", 12.0)), + "cpu", + ) + _, audio_timesteps = _make_schedule( + infer_steps, + float(config.get("audio_flow_shift", 3.0)), + "cpu", + ) + + entries = [] + for profile in profiles: + for step, (video_timestep, audio_timestep) in enumerate(zip(video_timesteps.tolist(), audio_timesteps.tolist())): + values = [video_timestep, audio_timestep] + if profile in {"conditioned", "ref2av_video", "ref2av_video_audio"}: + values.append(max(video_timestep, KEYFRAME_NOISE_AUG)) + if profile == "ref2av_video_audio": + values.append(CONDITION_AUDIO_TIMESTEP) + unique = torch.unique(torch.tensor(values, dtype=torch.float32), sorted=True) + entries.append( + { + "name": f"{profile}_step_{step:03d}", + "timestep_bits": _float32_bits(unique.tolist()), + } + ) + return entries + + +def _build_spec(config) -> dict: + if not config.get("use_adaln_cache", False): + raise ValueError("Building or loading an AdaLN cache requires use_adaln_cache=true") + validate_adaln_cache_config(config) + profiles = _selected_profiles(config) + return { + "infer_steps": int(config["infer_steps"]), + "video_flow_shift": float(config.get("video_flow_shift", 12.0)), + "audio_flow_shift": float(config.get("audio_flow_shift", 3.0)), + "num_layers": int(config.get("num_layers", 50)), + "hidden_size": int(config.get("hidden_size", 5376)), + "freq_dim": int(config.get("freq_dim", 256)), + "entries": _cache_entries(config, profiles), + } + + +def _cache_path(config) -> Path: + cache_name = "ref2av" if config["task"] == "ref2av" else "fl2av" + infer_steps = int(config["infer_steps"]) + return _cache_root(config) / "minimax_h3" / f"{cache_name}_{infer_steps:02d}steps" + + +def _expected_table_shape(spec: dict, entry: dict) -> tuple[int, int]: + return len(entry["timestep_bits"]) * 3, 6 * spec["hidden_size"] + + +def _expected_norm_out_shape(spec: dict, entry: dict) -> tuple[int, int]: + return len(entry["timestep_bits"]), 2 * spec["hidden_size"] + + +def _norm_out_key(entry: dict) -> str: + return f"norm_out.{entry['name']}" + + +def _block_key(block_index: int, entry: dict) -> str: + return f"block_{block_index:03d}.{entry['name']}" + + +def _validate_cache(cache_path: Path, spec: dict) -> bool: + manifest_path = cache_path / "manifest.json" + if not manifest_path.is_file(): + return False + try: + with manifest_path.open(encoding="utf-8") as handle: + manifest = json.load(handle) + if manifest != spec: + return False + expected_keys = {_norm_out_key(entry) for entry in spec["entries"]} + expected_keys.update(_block_key(block_index, entry) for block_index in range(spec["num_layers"]) for entry in spec["entries"]) + with safe_open(cache_path / "adaln_cache.safetensors", framework="pt", device="cpu") as source: + if set(source.keys()) != expected_keys: + return False + for entry in spec["entries"]: + tensor = source.get_slice(_norm_out_key(entry)) + if tuple(tensor.get_shape()) != _expected_norm_out_shape(spec, entry) or str(tensor.get_dtype()) != "BF16": + return False + for block_index in range(spec["num_layers"]): + for entry in spec["entries"]: + tensor = source.get_slice(_block_key(block_index, entry)) + if tuple(tensor.get_shape()) != _expected_table_shape(spec, entry) or str(tensor.get_dtype()) != "BF16": + return False + except (KeyError, OSError, RuntimeError, SafetensorError, TypeError, ValueError): + return False + return True + + +def load_persistent_adaln_cache( + config, + device, +) -> tuple[ + dict[tuple[float, ...], list[torch.Tensor]], + dict[tuple[float, ...], torch.Tensor], +]: + """Load cached block AdaLN and final-norm modulation onto the device.""" + spec = _build_spec(config) + cache_path = _cache_path(config) + if not _validate_cache(cache_path, spec): + raise FileNotFoundError(f"MiniMax-H3 AdaLN cache is missing or invalid: {cache_path}. Generate it first with tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py.") + + logger.info("========== Loading MiniMax-H3 AdaLN cache from {} ==========", cache_path) + keys = [tuple(_timesteps_from_bits(entry["timestep_bits"]).tolist()) for entry in spec["entries"]] + cache = {key: [None] * spec["num_layers"] for key in keys} + norm_out_cache = {} + with safe_open(cache_path / "adaln_cache.safetensors", framework="pt", device=str(device)) as source: + for entry, key in zip(spec["entries"], keys): + norm_out_cache[key] = source.get_tensor(_norm_out_key(entry)) + for block_index in range(spec["num_layers"]): + for entry, key in zip(spec["entries"], keys): + cache[key][block_index] = source.get_tensor(_block_key(block_index, entry)) + logger.success("========== MiniMax-H3 AdaLN cache loaded from {} ==========", cache_path) + return cache, norm_out_cache diff --git a/lightx2v/models/networks/minimax_h3/infer/module_io.py b/lightx2v/models/networks/minimax_h3/infer/module_io.py index d3b97da56..f19ca11d9 100644 --- a/lightx2v/models/networks/minimax_h3/infer/module_io.py +++ b/lightx2v/models/networks/minimax_h3/infer/module_io.py @@ -17,13 +17,14 @@ class MiniMaxH3SequenceParallelState: @dataclass class MiniMaxH3PreInferOutput: hidden_states: torch.Tensor - temb: torch.Tensor + temb: torch.Tensor | None timestep_indices: torch.Tensor adaln_indices: torch.Tensor rotary_emb: tuple[torch.Tensor, torch.Tensor] video_indices: torch.Tensor audio_indices: torch.Tensor text_indices: torch.Tensor + norm_out_modulation: torch.Tensor | None = None sequence_parallel_state: MiniMaxH3SequenceParallelState | None = None diff --git a/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py index 2093c6a6d..ad798fc76 100644 --- a/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py @@ -26,25 +26,8 @@ def get_compile_block_key(self, block_idx, block): # model offload return super().get_compile_block_key(block_idx, block) - def _prefetch_weights_without_adaln(self, block_index, blocks): - with torch_device_module.stream(self.offload_manager.cuda_load_stream): - if hasattr(self.offload_manager, "cpu_buffers"): - source_block = self.offload_manager.cpu_buffers[0] - else: - source_block = blocks[block_index] - block_state_dict = source_block.state_dict() - weights_without_adaln = {} - for name, tensor in block_state_dict.items(): - if ".adaln_proj." not in name: - weights_without_adaln[name] = tensor - self.offload_manager.cuda_buffers[1].load_state_dict(weights_without_adaln, block_index) - def infer_with_blocks_offload(self, blocks, hidden_states, pre_infer_out): num_blocks = len(blocks) - if self.use_adaln_cache and not self._adaln_cache_hit: - # The previous forward may have prefetched block 0 without AdaLN. - # Reload the full block when the current timestep misses. - self.offload_manager.need_init_first_buffer = True current_stream = torch_device_module.current_stream() self.offload_manager.compute_stream.wait_stream(current_stream) @@ -53,10 +36,7 @@ def infer_with_blocks_offload(self, blocks, hidden_states, pre_infer_out): self.offload_manager.init_first_buffer(blocks) next_block_index = (block_index + 1) % num_blocks - if self.use_adaln_cache and self._adaln_cache_hit: - self._prefetch_weights_without_adaln(next_block_index, blocks) - else: - self.offload_manager.prefetch_weights(next_block_index, blocks) + self.offload_manager.prefetch_weights(next_block_index, blocks) block = self.offload_manager.cuda_buffers[0] self.block_idx = block_index if AI_DEVICE == "xpu": diff --git a/lightx2v/models/networks/minimax_h3/infer/post_infer.py b/lightx2v/models/networks/minimax_h3/infer/post_infer.py index 5b88dc639..9143df184 100644 --- a/lightx2v/models/networks/minimax_h3/infer/post_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/post_infer.py @@ -12,7 +12,14 @@ def set_scheduler(self, scheduler): self.scheduler = scheduler def infer(self, weights, hidden_states, pre_infer_out): - shift, scale = weights.norm_out_linear.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())).chunk(2, dim=-1) + modulation = pre_infer_out.norm_out_modulation + if modulation is None: + # ADALN CACHE SYNC: The offline builder persists this exact + # norm_out.linear result. Mirror changes there and regenerate caches. + if pre_infer_out.temb is None: + raise RuntimeError("MiniMax-H3 final-norm modulation is missing") + modulation = weights.norm_out_linear.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())) + shift, scale = modulation.chunk(2, dim=-1) indices = pre_infer_out.timestep_indices hidden_states = weights.norm_out.apply(hidden_states) hidden_states = hidden_states * (1.0 + scale.index_select(0, indices)) diff --git a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py index 2ccd86806..f3ae92323 100644 --- a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py @@ -37,6 +37,7 @@ def __init__(self, config): self.rope_freq_dim = int(config.get("rope_freq_dim", 16)) self.rope_theta = float(config.get("rope_theta", 10000.0)) self.freq_dim = int(config.get("freq_dim", 256)) + self.use_adaln_cache = bool(config.get("use_adaln_cache", False)) def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -107,8 +108,13 @@ def infer(self, weights, prompt_embeds): hidden_states.index_copy_(0, layout.audio_indices, audio_embeds) hidden_states.index_copy_(0, layout.video_indices, video_embeds) - temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim) - temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float()))) + temb = None + if not self.use_adaln_cache: + # ADALN CACHE SYNC: Any change to this time-MLP sequence, activation, + # or dtype must also be made in the offline AdaLN cache builder and + # followed by regenerating the cache when cached values can change. + temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim) + temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float()))) timestep_indices = self.scheduler.timestep_indices adaln_indices = timestep_indices * 3 + layout.token_tags.clamp(min=0) diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index f7366425f..299d4c3ad 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -3,7 +3,9 @@ import torch.nn.functional as F from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer +from lightx2v.models.networks.minimax_h3.adaln_cache import load_persistent_adaln_cache from lightx2v.utils.envs import GET_DTYPE +from lightx2v_platform.base.global_var import AI_DEVICE class MiniMaxH3TransformerInfer(BaseTransformerInfer): @@ -36,11 +38,17 @@ def __init__(self, config): self.seq_p_group = None self.infer_func = self.infer_without_offload self.use_adaln_cache = bool(config.get("use_adaln_cache", False)) - self._adaln_cache = {} self._current_adaln_tables = None self._adaln_cache_hit = False + if self.use_adaln_cache: + self._adaln_cache, self._norm_out_cache = load_persistent_adaln_cache(config, self._cache_device()) self.init_compile(config) + @staticmethod + def _cache_device(): + device_module = getattr(torch, AI_DEVICE) + return torch.device(AI_DEVICE, device_module.current_device()) + def _gather_tp_last_dim(self, tensor): if self.tp_size == 1: return tensor @@ -128,38 +136,42 @@ def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None): return hidden_states def _compute_adaln_table(self, weights, pre_infer_out): + # ADALN CACHE SYNC: This projection is reproduced by the offline builder. + # Keep the offline AdaLN cache builder and its tests aligned with changes + # to activation placement, dtype, projection, gather, or reshape. # Activation is evaluated in fp32, then cast to the inference dtype # immediately before the (possibly quantized) AdaLN projection. + if pre_infer_out.temb is None: + raise RuntimeError("MiniMax-H3 timestep embedding is missing") modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(self.infer_dtype)) modulation = self._gather_tp_last_dim(modulation) return modulation.view(-1, 6 * self.hidden_size) def _clear_adaln_cache(self): - self._adaln_cache.clear() + # Online mode has no reusable cache; it only creates a temporary table + # for the current block invocation in _compute_adaln_table(). self._current_adaln_tables = None self._adaln_cache_hit = False - def _prepare_adaln_cache(self): + def _prepare_adaln_cache(self, pre_infer_out): current_timesteps = tuple(self.scheduler.unique_timesteps_cpu.tolist()) cached_tables = self._adaln_cache.get(current_timesteps) - if cached_tables is not None: + norm_out_modulation = self._norm_out_cache.get(current_timesteps) + if cached_tables is not None and norm_out_modulation is not None: self._current_adaln_tables = cached_tables + pre_infer_out.norm_out_modulation = norm_out_modulation self._adaln_cache_hit = True else: - self._current_adaln_tables = [] - self._adaln_cache[current_timesteps] = self._current_adaln_tables - self._adaln_cache_hit = False + raise KeyError(f"Persistent MiniMax-H3 AdaLN cache has no entry for timesteps {current_timesteps}") - def _get_or_build_adaln(self, block_index, weights, pre_infer_out): + def _get_cached_adaln(self, block_index): if self._adaln_cache_hit: return self._current_adaln_tables[block_index] - adaln_table = self._compute_adaln_table(weights, pre_infer_out) - self._current_adaln_tables.append(adaln_table) - return adaln_table + raise RuntimeError("Persistent MiniMax-H3 AdaLN cache was not prepared") def run_block(self, block_idx, block, hidden_states, pre_infer_out): if self.use_adaln_cache: - adaln_table = self._get_or_build_adaln(block_idx, block, pre_infer_out) + adaln_table = self._get_cached_adaln(block_idx) return super().run_block(block_idx, block, hidden_states, pre_infer_out, adaln_table) return super().run_block(block_idx, block, hidden_states, pre_infer_out) @@ -171,5 +183,5 @@ def infer_without_offload(self, blocks, hidden_states, pre_infer_out): def infer(self, block_weights, pre_infer_out): if self.use_adaln_cache: - self._prepare_adaln_cache() + self._prepare_adaln_cache(pre_infer_out) return self.infer_func(block_weights.blocks, pre_infer_out.hidden_states, pre_infer_out) diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 19ca82476..94d02ed4c 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -8,6 +8,7 @@ from safetensors import safe_open from lightx2v.models.networks.base_model import BaseTransformerModel +from lightx2v.models.networks.minimax_h3.adaln_cache import validate_adaln_cache_config from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3SequenceParallelState from lightx2v.models.networks.minimax_h3.infer.offload import MiniMaxH3OffloadTransformerInfer from lightx2v.models.networks.minimax_h3.infer.post_infer import MiniMaxH3PostInfer @@ -48,6 +49,31 @@ class MiniMaxH3Model(BaseTransformerModel): def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0, lora_alpha=None): self.lora_alpha = lora_alpha + self.use_adaln_cache = bool(config.get("use_adaln_cache", False)) + if config.get("cpu_offload", False) and not self.use_adaln_cache: + separator = "=" * 88 + message = ( + f"\n{separator}\n" + "MINIMAX-H3 CPU OFFLOAD CONFIGURATION ERROR\n" + "cpu_offload=true requires use_adaln_cache=true.\n" + "\nACTION REQUIRED\n" + "Enable the AdaLN cache and set its root in the inference JSON config:\n" + ' "use_adaln_cache": true,\n' + ' "adaln_cache_dir": "~/.cache/lightx2v/adaln"\n' + "adaln_cache_dir can be any custom cache root; the path above is the recommended default.\n" + "\nBefore inference, set lightx2v_path, model_path, --config_json, and " + "--task in tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh.\n" + "Then generate the cache from the repository root with:\n" + " bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh\n" + "Set --task fl2av in that script for t2av/i2av/l2av/fl2av, or " + "--task ref2av for ref2av. The generation script and inference must " + f"use the same JSON config.\n{separator}" + ) + logger.error(message) + raise ValueError(message) + if self.use_adaln_cache: + validate_adaln_cache_config(config) + self.remove_keys = [".adaln_proj.", "time_embedder.", "norm_out.linear."] self.block_offload = config.get("cpu_offload", False) and config.get("offload_granularity", "model") == "block" # Model offload moves pre/blocks/post together. Pre/post residency only applies # to block offload and is ignored otherwise. @@ -297,8 +323,9 @@ def _validate_tensor_parallel_config(self): checks = { "num_attention_heads": int(self.config.get("num_attention_heads", 56)), "ffn_hidden_size": int(self.config.get("ffn_hidden_size", 14336)), - "adaln_output_size": 18 * int(self.config.get("hidden_size", 5376)), } + if not self.use_adaln_cache: + checks["adaln_output_size"] = 18 * int(self.config.get("hidden_size", 5376)) invalid = {name: value for name, value in checks.items() if value % self.tp_size} if invalid: details = ", ".join(f"{name}={value}" for name, value in invalid.items()) diff --git a/lightx2v/models/networks/minimax_h3/packing.py b/lightx2v/models/networks/minimax_h3/packing.py index b0a2dc950..9024a3382 100644 --- a/lightx2v/models/networks/minimax_h3/packing.py +++ b/lightx2v/models/networks/minimax_h3/packing.py @@ -31,6 +31,7 @@ PIXEL_MEAN = (0.485, 0.456, 0.406) PIXEL_STD = (0.229, 0.224, 0.225) KEYFRAME_NOISE_AUG = 0.999 +CONDITION_AUDIO_TIMESTEP = 1.0 KEYFRAME_ENCODE_SEED = 42 _ROPE_FRAME_RESCALE = 5.0 / 3.0 @@ -287,7 +288,7 @@ def build_row_timesteps( video_timestep: float, audio_timestep: float, condition_video_timestep: float | None = None, - condition_audio_timestep: float = 1.0, + condition_audio_timestep: float = CONDITION_AUDIO_TIMESTEP, ) -> tuple[torch.Tensor, torch.Tensor]: row_timesteps = torch.full((layout.sequence_length,), video_timestep, dtype=torch.float32) if condition_video_timestep is None: diff --git a/lightx2v/models/networks/minimax_h3/weights/post_weights.py b/lightx2v/models/networks/minimax_h3/weights/post_weights.py index 43c7b0f90..322e7b8a9 100644 --- a/lightx2v/models/networks/minimax_h3/weights/post_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/post_weights.py @@ -13,10 +13,13 @@ def __init__(self, config): "norm_out", _rms(config, "norm_out.norm.weight", eps=float(config.get("final_norm_eps", 1e-5))), ) - self.add_module( - "norm_out_linear", - MM_WEIGHT_REGISTER["Default"]("norm_out.linear.weight", "norm_out.linear.bias"), - ) + if not config.get("use_adaln_cache", False): + # ADALN CACHE SYNC: The offline builder reads this key and persists + # its output; update the offline builder if its definition changes. + self.add_module( + "norm_out_linear", + MM_WEIGHT_REGISTER["Default"]("norm_out.linear.weight", "norm_out.linear.bias"), + ) self.add_module( "proj_out", MM_WEIGHT_REGISTER["Default-ForceFp32"]("proj_out.weight", "proj_out.bias"), diff --git a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py b/lightx2v/models/networks/minimax_h3/weights/pre_weights.py index 974dab411..76b46f91c 100644 --- a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/pre_weights.py @@ -84,8 +84,11 @@ def __init__(self, config): self.add_module("proj_in", _linear("proj_in", bias=True, force_fp32=True)) self.add_module("audio_proj_in", _linear("audio_proj_in", bias=True, force_fp32=True)) self.add_module("context_embedder", _linear("context_embedder", bias=True)) - self.add_module("time_linear_1", _linear("time_embedder.linear_1", bias=True, force_fp32=True)) - self.add_module("time_linear_2", _linear("time_embedder.linear_2", bias=True, force_fp32=True)) + if not config.get("use_adaln_cache", False): + # ADALN CACHE SYNC: The offline builder reads these keys and mirrors + # their FP32 semantics; update the offline builder if either changes. + self.add_module("time_linear_1", _linear("time_embedder.linear_1", bias=True, force_fp32=True)) + self.add_module("time_linear_2", _linear("time_embedder.linear_2", bias=True, force_fp32=True)) self.add_module( "refiner_blocks", WeightModuleList([MiniMaxH3TokenRefinerBlockWeights(i, config) for i in range(int(config.get("num_refiner_layers", 2)))]), diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..d7fa24758 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -121,9 +121,12 @@ def __init__(self, index, config, create_cuda_buffer=False): ), ) self.add_module("ff", MiniMaxH3FeedForwardWeights(f"{prefix}.ff", config, create_cuda_buffer)) - # AdaLN is the largest per-block projection in H3. Its output is - # column-sharded here and gathered once per block before modulation. - self.add_module("adaln", _linear(config, f"{prefix}.adaln_proj.linear", bias=True, create_cuda_buffer=create_cuda_buffer, tp_split="col")) + if not config.get("use_adaln_cache", False): + # ADALN CACHE SYNC: The offline builder reads this key and mirrors + # the unquantized projection; update the offline builder if it changes. + # AdaLN is the largest per-block projection in H3. Its output is + # column-sharded here and gathered once per block before modulation. + self.add_module("adaln", _linear(config, f"{prefix}.adaln_proj.linear", bias=True, create_cuda_buffer=create_cuda_buffer, tp_split="col")) class MiniMaxH3TransformerWeights(WeightModule): diff --git a/scripts/minimax_h3/README.md b/scripts/minimax_h3/README.md index 6b5a54519..344662cb9 100644 --- a/scripts/minimax_h3/README.md +++ b/scripts/minimax_h3/README.md @@ -31,6 +31,16 @@ Each transformer directory needs its `config.json`, weight index, and checkpoint The base transformer can serve all four base tasks without reloading. Reference generation uses a separate transformer and service. These checkpoints are CFG-distilled: do not send `negative_prompt`, including an empty string. +## AdaLN cache + +MiniMax-H3 CPU offload requires `use_adaln_cache: true`. Before starting inference, generate the persistent cache with the same model, config, inference-step count, and flow shifts that inference will use: + +```bash +bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +``` + +Set `lightx2v_path`, `model_path`, `--config_json`, and `--task` in the script before running it. `--task fl2av` creates both base-transformer profiles and serves `t2av`, `i2av`, `l2av`, and `fl2av`. Run the script separately with `--task ref2av` to build the reference-transformer cache. Every JSON config with `use_adaln_cache: true` must explicitly set `adaln_cache_dir`; the bundled configs use `~/.cache/lightx2v/adaln`. Final directory names include the cache group and inference-step count, such as `minimax_h3/fl2av_29steps`, `minimax_h3/fl2av_04steps`, or `minimax_h3/ref2av_29steps`. Each `manifest.json` stores the minimal cache specification, which inference compares directly with its expected specification. Offline generation and inference read this same JSON setting. Cache generation refuses to overwrite an existing target directory. Inference loads a matching cache strictly and does not fall back to online AdaLN computation. + ## Offline inference The five task scripts share `configs/minimax_h3/minimax_h3.json`: one GPU, BF16 weights, model CPU offload, and 124 frames at `[height, width] = [544, 960]`. The script's `--task` selects the transformer and input handling; the JSON filename does not select a task. diff --git a/scripts/minimax_h3/README_zh.md b/scripts/minimax_h3/README_zh.md index eebd86a14..d8a2dec24 100644 --- a/scripts/minimax_h3/README_zh.md +++ b/scripts/minimax_h3/README_zh.md @@ -31,6 +31,16 @@ MiniMax-H3/ 基础 transformer 加载一次即可处理前四种任务。参考生成使用独立的 transformer 和服务。这些权重已经完成 CFG 蒸馏，请勿传入 `negative_prompt`，包括空字符串。 +## AdaLN 缓存 + +MiniMax-H3 开启 CPU offload 时强制要求 `use_adaln_cache: true`。开始推理前，需要使用与推理一致的模型、配置、推理步数和 flow shift 生成持久化缓存： + +```bash +bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +``` + +运行前需要在脚本中设置 `lightx2v_path`、`model_path`、`--config_json` 和 `--task`。使用 `--task fl2av` 会生成两套基础 transformer profile，可供 `t2av`、`i2av`、`l2av` 和 `fl2av` 共用。参考任务需要另外使用 `--task ref2av` 生成 `transformer_ref` 缓存。所有设置了 `use_adaln_cache: true` 的 JSON 配置都必须显式设置 `adaln_cache_dir`；仓库自带配置统一写为 `~/.cache/lightx2v/adaln`。最终目录名包含缓存任务组和推理步数，例如 `minimax_h3/fl2av_29steps`、`minimax_h3/fl2av_04steps` 或 `minimax_h3/ref2av_29steps`。各自的 `manifest.json` 只保存必要的缓存规格，推理时直接与当前预期规格比较。离线生成和推理读取同一个 JSON 配置项。缓存生成不会覆盖已经存在的目标目录。推理会严格读取匹配缓存，不会回退到在线 AdaLN 计算。 + ## 离线推理 五种任务脚本共用 `configs/minimax_h3/minimax_h3.json`：单 GPU、BF16 权重、模型级 CPU 卸载，默认输出 124 帧，`[高度, 宽度] = [544, 960]`。脚本中的 `--task` 决定加载哪组 transformer 以及如何处理输入，JSON 文件名不决定任务。 diff --git a/scripts/minimax_h3/run_minimax_h3_fl2av.sh b/scripts/minimax_h3/run_minimax_h3_fl2av.sh index 1e5c9a54b..d3eb9833f 100755 --- a/scripts/minimax_h3/run_minimax_h3_fl2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_fl2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/run_minimax_h3_i2av.sh b/scripts/minimax_h3/run_minimax_h3_i2av.sh index 980522c72..27f201243 100755 --- a/scripts/minimax_h3/run_minimax_h3_i2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_i2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/run_minimax_h3_l2av.sh b/scripts/minimax_h3/run_minimax_h3_l2av.sh index 0e001236c..74aac3f74 100755 --- a/scripts/minimax_h3/run_minimax_h3_l2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_l2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/run_minimax_h3_ref2av.sh b/scripts/minimax_h3/run_minimax_h3_ref2av.sh index 5bbe48bc2..0a98f0374 100755 --- a/scripts/minimax_h3/run_minimax_h3_ref2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_ref2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av.sh b/scripts/minimax_h3/run_minimax_h3_t2av.sh index 904dee980..ff57c257a 100755 --- a/scripts/minimax_h3/run_minimax_h3_t2av.sh +++ b/scripts/minimax_h3/run_minimax_h3_t2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh b/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh index feb5b1326..69f3cbf93 100644 --- a/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh +++ b/scripts/minimax_h3/run_minimax_h3_t2av_parallel.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # set path firstly lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/server/start_server.sh b/scripts/minimax_h3/server/start_server.sh index 25c309a42..e90c3baac 100644 --- a/scripts/minimax_h3/server/start_server.sh +++ b/scripts/minimax_h3/server/start_server.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/minimax_h3/server/start_server_ref2av.sh b/scripts/minimax_h3/server/start_server_ref2av.sh index 33a0191ad..75e754e1d 100644 --- a/scripts/minimax_h3/server/start_server_ref2av.sh +++ b/scripts/minimax_h3/server/start_server_ref2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + lightx2v_path=/path/to/LightX2V model_path=/path/to/MiniMax-H3 diff --git a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_dmd_lora_4step_sp_tp.sh b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_dmd_lora_4step_sp_tp.sh index e3757fa63..384600aeb 100755 --- a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_dmd_lora_4step_sp_tp.sh +++ b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_dmd_lora_4step_sp_tp.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_int8_sp_tp.sh b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_int8_sp_tp.sh index 1fbd08d30..c88ba6142 100755 --- a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_int8_sp_tp.sh +++ b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_int8_sp_tp.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp.sh b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp.sh index 46b0b301a..31d6f4e05 100755 --- a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp.sh +++ b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp_cpu_offload.sh b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp_cpu_offload.sh index 6df51933c..718cfd022 100755 --- a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp_cpu_offload.sh +++ b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_sp_tp_cpu_offload.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_tp.sh b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_tp.sh index 553fb9cf0..f1b5e80f2 100755 --- a/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_tp.sh +++ b/scripts/platforms/intel_xpu/dist_infer/run_minimax_h3_t2av_tp.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh index ecb495b7b..f360bc82d 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail # Usage: diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_t2av.sh b/scripts/platforms/intel_xpu/run_minimax_h3_t2av.sh index 42c84b30c..623d34927 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_t2av.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_t2av.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_dmd_lora_4step.sh b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_dmd_lora_4step.sh index eefa5a45a..415b1ec4d 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_dmd_lora_4step.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_dmd_lora_4step.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_fp8.sh b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_fp8.sh index 21ff32571..150c2d9a7 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_fp8.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_fp8.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_int8.sh b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_int8.sh index 03e7afaee..0abc31440 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_t2av_int8.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_t2av_int8.sh @@ -1,4 +1,13 @@ #!/usr/bin/env bash + +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. set -euo pipefail SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) diff --git a/scripts/platforms/metax/run_minimax_h3_t2av_tp1_block_offload.sh b/scripts/platforms/metax/run_minimax_h3_t2av_tp1_block_offload.sh index 42555d704..a9d8bd9d0 100755 --- a/scripts/platforms/metax/run_minimax_h3_t2av_tp1_block_offload.sh +++ b/scripts/platforms/metax/run_minimax_h3_t2av_tp1_block_offload.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + lightx2v_path=${LIGHTX2V_PATH:-/data/LightX2V} model_path=${MINIMAX_H3_MODEL_PATH:-/data/models/MiniMax-H3} save_result_path=${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp1_block_offload.mp4} diff --git a/scripts/platforms/metax/run_minimax_h3_t2av_tp_sp.sh b/scripts/platforms/metax/run_minimax_h3_t2av_tp_sp.sh index a2a7772f8..85c34c345 100755 --- a/scripts/platforms/metax/run_minimax_h3_t2av_tp_sp.sh +++ b/scripts/platforms/metax/run_minimax_h3_t2av_tp_sp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + lightx2v_path=${LIGHTX2V_PATH:-/data/LightX2V} model_path=${MINIMAX_H3_MODEL_PATH:-/data/models/MiniMax-H3} save_result_path=${SAVE_RESULT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp_sp.mp4} diff --git a/scripts/platforms/mlu/run_minimax_h3_t2av_sp.sh b/scripts/platforms/mlu/run_minimax_h3_t2av_sp.sh index d3a9eb27c..9ef016c1b 100644 --- a/scripts/platforms/mlu/run_minimax_h3_t2av_sp.sh +++ b/scripts/platforms/mlu/run_minimax_h3_t2av_sp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mlu/run_minimax_h3_t2av_tp.sh b/scripts/platforms/mlu/run_minimax_h3_t2av_tp.sh index 2ebc9ccde..c424a9a88 100644 --- a/scripts/platforms/mlu/run_minimax_h3_t2av_tp.sh +++ b/scripts/platforms/mlu/run_minimax_h3_t2av_tp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mlu/run_minimax_h3_t2av_tp_sp.sh b/scripts/platforms/mlu/run_minimax_h3_t2av_tp_sp.sh index fb1aef62b..6931f28ce 100644 --- a/scripts/platforms/mlu/run_minimax_h3_t2av_tp_sp.sh +++ b/scripts/platforms/mlu/run_minimax_h3_t2av_tp_sp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av.sh index 1eb0d6992..290372fdb 100644 --- a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av.sh +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_fp8.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_fp8.sh index aacd78a3a..477ecfad7 100644 --- a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_fp8.sh +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_fp8.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp.sh index 6bacb0381..da4f6dc85 100644 --- a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp.sh +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_fp8.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_fp8.sh index 30152fda6..c4e2e4f50 100644 --- a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_fp8.sh +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_fp8.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_sp.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_sp.sh index 736e01631..55179b11e 100644 --- a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_sp.sh +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp_sp.sh @@ -1,5 +1,14 @@ #!/bin/bash +# AdaLN cache setup: +# If the inference JSON config enables "use_adaln_cache": true, generate the cache before inference: +# 1. Set lightx2v_path, model_path, --config_json, and --task in +# tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh. +# 2. Use --task fl2av for t2av/i2av/l2av/fl2av, or --task ref2av for ref2av. +# 3. From the repository root, run: +# bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +# Cache generation and inference must use the same JSON config and adaln_cache_dir. + # System management interface: mthreads-gmi # set path firstly diff --git a/tools/cache_minimax_h3_adaln/builder.py b/tools/cache_minimax_h3_adaln/builder.py new file mode 100644 index 000000000..5aef69535 --- /dev/null +++ b/tools/cache_minimax_h3_adaln/builder.py @@ -0,0 +1,176 @@ +"""Build a MiniMax-H3 AdaLN cache without loading the full model. + +ADALN CACHE SYNC CONTRACT: +The calculations in ``_build_cache`` intentionally mirror the online BF16 +path in ``infer/pre_infer.py``, ``infer/transformer_infer.py``, and +``infer/post_infer.py``. If timestep embedding, time-MLP activation/dtype, +AdaLN projection/reshape, final-norm modulation, or their checkpoint keys are +changed online, update this builder and the cache tests in the same change. +Regenerate the cache whenever its values can change. +""" + +import json +import os +import shutil +import tempfile +from pathlib import Path + +import torch +import torch.nn.functional as F +from loguru import logger +from safetensors import safe_open +from safetensors.torch import save_file + +from lightx2v.models.networks.minimax_h3.adaln_cache import ( + _block_key, + _build_spec, + _cache_path, + _expected_table_shape, + _norm_out_key, + _timesteps_from_bits, + _validate_cache, +) +from lightx2v.models.networks.minimax_h3.infer.pre_infer import timestep_embedding +from lightx2v_platform.base.global_var import AI_DEVICE + +torch_device_module = getattr(torch, AI_DEVICE) + + +def _checkpoint_files(config) -> list[Path]: + checkpoint = Path(config["dit_original_ckpt"]).expanduser().resolve() + files = sorted(checkpoint.glob("*.safetensors")) if checkpoint.is_dir() else [checkpoint] + if not files or any(not path.is_file() for path in files): + raise FileNotFoundError(f"MiniMax-H3 safetensors checkpoint not found: {checkpoint}") + return files + + +class _CheckpointTensors: + """Read individual tensors without materializing the whole checkpoint.""" + + def __init__(self, files: list[Path]): + self.files = files + self.locations = self._find_locations() + + def _find_locations(self) -> dict[str, Path]: + directory = self.files[0].parent + index_files = sorted(directory.glob("*.safetensors.index.json")) + if index_files: + with index_files[0].open(encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + return {name: directory / filename for name, filename in weight_map.items()} + + locations = {} + for path in self.files: + with safe_open(path, framework="pt", device="cpu") as source: + locations.update({name: path for name in source.keys()}) + return locations + + def get(self, name: str) -> torch.Tensor: + path = self.locations.get(name) + if path is None: + raise KeyError(f"MiniMax-H3 checkpoint tensor is missing: {name}") + with safe_open(path, framework="pt", device="cpu") as source: + return source.get_tensor(name) + + +def _linear( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, +) -> torch.Tensor: + # ADALN CACHE SYNC: Match the transpose, bias, output dtype, and operation + # order of the online Default MMWeight.apply path. Do not change this + # independently of MiniMax-H3's online modulation projections. + output = torch.empty( + (input_tensor.shape[0], weight.shape[0]), + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + return torch.addmm(bias, input_tensor, weight.t(), out=output) + + +def _empty_device_cache() -> None: + if hasattr(torch_device_module, "empty_cache"): + torch_device_module.empty_cache() + + +def _build_cache(spec: dict, cache_path: Path, checkpoint_files: list[Path]) -> None: + stage_path = Path(tempfile.mkdtemp(prefix=".building-", dir=cache_path.parent)) + try: + checkpoint = _CheckpointTensors(checkpoint_files) + with torch.inference_mode(): + # ADALN CACHE SYNC: Keep activation placement and casts aligned with + # the three online infer modules named in this file's contract. + time_weight_1 = checkpoint.get("time_embedder.linear_1.weight").to(AI_DEVICE) + time_bias_1 = checkpoint.get("time_embedder.linear_1.bias").to(AI_DEVICE) + time_weight_2 = checkpoint.get("time_embedder.linear_2.weight").to(AI_DEVICE) + time_bias_2 = checkpoint.get("time_embedder.linear_2.bias").to(AI_DEVICE) + + adaln_inputs = {} + for entry in spec["entries"]: + timesteps = _timesteps_from_bits(entry["timestep_bits"], AI_DEVICE) + embedded = timestep_embedding(timesteps, spec["freq_dim"]) + temb = _linear(embedded.float(), time_weight_1, time_bias_1) + temb = _linear(F.silu(temb), time_weight_2, time_bias_2) + adaln_inputs[entry["name"]] = F.silu(temb).to(torch.bfloat16) + + del time_weight_1, time_bias_1, time_weight_2, time_bias_2 + _empty_device_cache() + + norm_out_weight = checkpoint.get("norm_out.linear.weight").to(AI_DEVICE) + norm_out_bias = checkpoint.get("norm_out.linear.bias").to(AI_DEVICE) + cache_tables = { + _norm_out_key(entry): _linear( + adaln_inputs[entry["name"]], + norm_out_weight, + norm_out_bias, + ) + .to("cpu") + .contiguous() + for entry in spec["entries"] + } + del norm_out_weight, norm_out_bias + _empty_device_cache() + + # Only one full, unsharded block projection is resident at a time. + for block_index in range(spec["num_layers"]): + prefix = f"transformer_blocks.{block_index}.adaln_proj.linear" + weight = checkpoint.get(f"{prefix}.weight").to(AI_DEVICE) + bias = checkpoint.get(f"{prefix}.bias").to(AI_DEVICE) + for entry in spec["entries"]: + modulation = _linear(adaln_inputs[entry["name"]], weight, bias) + table = modulation.view(*_expected_table_shape(spec, entry)) + cache_tables[_block_key(block_index, entry)] = table.to("cpu").contiguous() + del weight, bias, modulation, table + _empty_device_cache() + logger.info( + "Built MiniMax-H3 AdaLN cache block {}/{}", + block_index + 1, + spec["num_layers"], + ) + + save_file(cache_tables, stage_path / "adaln_cache.safetensors") + + with (stage_path / "manifest.json").open("w", encoding="utf-8") as handle: + json.dump(spec, handle, ensure_ascii=False, indent=2, sort_keys=True) + os.replace(stage_path, cache_path) + finally: + if stage_path.exists(): + shutil.rmtree(stage_path) + + +def build_persistent_adaln_cache(config) -> Path: + """Build the cache in this process for later inference processes to load.""" + spec = _build_spec(config) + cache_path = _cache_path(config) + cache_path.parent.mkdir(parents=True, exist_ok=True) + if cache_path.exists(): + raise FileExistsError(f"MiniMax-H3 AdaLN cache path already exists: {cache_path}") + checkpoint_files = _checkpoint_files(config) + logger.info("Building MiniMax-H3 AdaLN cache on {}: {}", AI_DEVICE, cache_path) + _build_cache(spec, cache_path, checkpoint_files) + + if not _validate_cache(cache_path, spec): + raise RuntimeError(f"MiniMax-H3 AdaLN cache validation failed: {cache_path}") + logger.info("MiniMax-H3 AdaLN cache is ready: {}", cache_path) + return cache_path diff --git a/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py b/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py new file mode 100644 index 000000000..f5045a911 --- /dev/null +++ b/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Generate persistent MiniMax-H3 timestep modulation without loading the model.""" + +import argparse +import sys +from pathlib import Path + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +if str(_REPOSITORY_ROOT) not in sys.path: + sys.path.insert(0, str(_REPOSITORY_ROOT)) + +from loguru import logger # noqa: E402 + +from lightx2v.utils.set_config import build_startup_config # noqa: E402 +from tools.cache_minimax_h3_adaln.builder import ( # noqa: E402 + build_persistent_adaln_cache, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Build the on-disk AdaLN and final-norm cache required by MiniMax-H3 inference.", + ) + parser.add_argument("--model_path", required=True, help="MiniMax-H3 model root") + parser.add_argument("--config_json", required=True, help="Inference JSON config") + parser.add_argument( + "--task", + required=True, + choices=("fl2av", "ref2av"), + help="Cache the two base-transformer profiles for fl2av or the two reference-transformer profiles for ref2av", + ) + parser.set_defaults(model_cls="minimax_h3") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + config = build_startup_config( + { + "model_cls": args.model_cls, + "model_path": args.model_path, + "config_json": args.config_json, + "task": args.task, + } + ) + cache_path = build_persistent_adaln_cache(config) + logger.info("MiniMax-H3 AdaLN cache saved to {}", cache_path) + + +if __name__ == "__main__": + main() diff --git a/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh b/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh new file mode 100755 index 000000000..8f0b5f950 --- /dev/null +++ b/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/data/nvme1/yongyang/dan/LightX2V +model_path=/data/nvme1/models/MiniMaxAI/MiniMax-H3 + +# Select one platform. NVIDIA is enabled by default. + +# NVIDIA +export PLATFORM=cuda +export CUDA_VISIBLE_DEVICES=0 + +# Intel XPU +# export PLATFORM=intel_xpu +# export ZE_AFFINITY_MASK=0 + +# AMD ROCm +# export PLATFORM=amd_rocm +# export CUDA_VISIBLE_DEVICES=0 + +# MetaX +# export PLATFORM=metax_cuda +# export CUDA_VISIBLE_DEVICES=0 + +# Ascend NPU +# export PLATFORM=ascend_npu +# export ASCEND_RT_VISIBLE_DEVICES=0 + +# MThreads MUSA +# export PLATFORM=musa +# export MUSA_VISIBLE_DEVICES=0 + +# Cambricon MLU +# export PLATFORM=cambricon_mlu +# export MLU_VISIBLE_DEVICES=0 + +# Hygon DCU +# export PLATFORM=hygon_dcu +# export HIP_VISIBLE_DEVICES=0 + +# Enflame GCU +# export PLATFORM=enflame_gcu +# export ECCL_RAS_DISABLE=2 + +# Iluvatar CoreX +# export PLATFORM=iluvatar_cuda +# export CUDA_VISIBLE_DEVICES=0 + +# PPU +# export PLATFORM=ppu_cuda +# export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source "${lightx2v_path}/scripts/base/base.sh" + +# Supported tasks: fl2av, ref2av +python "${lightx2v_path}/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py" \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --task fl2av