From 7a83b80ec7ac94f7893ca96c0e749f74aa4444ff Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 4 Mar 2026 18:30:59 +0000 Subject: [PATCH 1/4] Update the ltx-2 Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 4 +++- examples/diffusers/quantization/pipeline_manager.py | 1 + examples/diffusers/quantization/quantize.py | 13 ++++++++++--- examples/diffusers/quantization/utils.py | 2 +- 4 files changed, 15 insertions(+), 5 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 17b010554..1a89a306f 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -126,7 +126,9 @@ python quantize.py \ --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \ --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \ --extra-param fp8transformer=true \ - --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt + --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \ + --hf-ckpt-dir ./LTX2-NVFP4/ \ + --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors ``` #### Important Parameters diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index 68ddd8791..6575e7e83 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -216,6 +216,7 @@ def _create_ltx2_pipeline(self) -> Any: fp8_quantization = params.pop("fp8_quantization", None) or params.pop( "fp8transformer", False ) + params.pop("merged_base_safetensor_path", None) if not checkpoint_path: raise ValueError("Missing required extra_param: checkpoint_path.") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index bfff207af..612357f6e 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -297,18 +297,25 @@ def restore_checkpoint(self) -> None: self.logger.info("Backbone checkpoints restored successfully") # TODO: should not do the any data type - def export_hf_ckpt(self, pipe: Any) -> None: + def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None: """ Export quantized model to HuggingFace checkpoint format. Args: pipe: Diffusion pipeline containing the quantized model + model_config: Model configuration (used to pass model-specific export kwargs) """ if not self.config.hf_ckpt_dir: return self.logger.info(f"Exporting HuggingFace checkpoint to {self.config.hf_ckpt_dir}") - export_hf_checkpoint(pipe, export_dir=self.config.hf_ckpt_dir) + kwargs: dict[str, Any] = {} + if model_config and model_config.model_type == ModelType.LTX2: + merged_path = model_config.extra_params.get("merged_base_safetensor_path") + if merged_path: + self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export") + kwargs["merged_base_safetensor_path"] = merged_path + export_hf_checkpoint(pipe, export_dir=self.config.hf_ckpt_dir, **kwargs) self.logger.info("HuggingFace checkpoint export completed successfully") @@ -601,7 +608,7 @@ def forward_loop(mod): quant_config.format, ) - export_manager.export_hf_ckpt(pipe) + export_manager.export_hf_ckpt(pipe, model_config) logger.info( f"Quantization process completed successfully! Time taken = {time.time() - s} seconds" diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index 21fcd87d0..c354b759c 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -84,7 +84,7 @@ def filter_func_flux_dev(name: str) -> bool: def filter_func_wan_video(name: str) -> bool: """Filter function specifically for LTX-Video models.""" pattern = re.compile( - r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.39|blocks.38).*" + r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.2\.|blocks.39|blocks.38|blocks.39).*" ) return pattern.match(name) is not None From 24030528961f3edd49b94f83ceccc85804872779 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 6 Mar 2026 00:55:07 +0000 Subject: [PATCH 2/4] Update Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 4 +++- examples/diffusers/quantization/utils.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 1a89a306f..7b50f4ee2 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -116,7 +116,9 @@ python quantize.py \ --hf-ckpt-dir ./hf_ckpt ``` -#### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 (torch checkpoint export) +#### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 + +This example produces three outputs: a PyTorch checkpoint (`--quantized-torch-ckpt-save-path`), a Hugging Face checkpoint (`--hf-ckpt-dir`), and a ComfyUI-compatible merged safetensor (`--extra-param merged_base_safetensor_path`). ```sh python quantize.py \ diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c354b759c..d89915da2 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -82,9 +82,9 @@ def filter_func_flux_dev(name: str) -> bool: def filter_func_wan_video(name: str) -> bool: - """Filter function specifically for LTX-Video models.""" + """Filter function specifically for WAN-Video models.""" pattern = re.compile( - r".*(patch_embedding|condition_embedder|proj_out|blocks.0\.|blocks.1\.|blocks.2\.|blocks.39|blocks.38|blocks.39).*" + r".*(patch_embedding|condition_embedder|proj_out|blocks\.(0|1|2|37|38|39)\.).*" ) return pattern.match(name) is not None From 5ebce973e2ded27bf36815a6276a53be718a3ec6 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 6 Mar 2026 05:08:07 +0000 Subject: [PATCH 3/4] Update the recipe Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index d89915da2..fb0231204 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -70,7 +70,7 @@ def check_conv_and_mha(backbone, if_fp4, quantize_mha): def filter_func_ltx_video(name: str) -> bool: """Filter function specifically for LTX-Video models.""" pattern = re.compile( - r".*(proj_in|time_embed|caption_projection|proj_out|patchify_proj|adaln_single).*" + r".*(proj_in|time_embed|caption_projection|proj_out|patchify_proj|adaln_single|transformer_blocks\.(0|1|2|46|47)\.).*" ) return pattern.match(name) is not None From 4e5c4ed1bb83b8136e95685b9669d66531342c92 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Sat, 7 Mar 2026 00:19:11 -0800 Subject: [PATCH 4/4] Update Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index fb0231204..f88c55a53 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -70,7 +70,7 @@ def check_conv_and_mha(backbone, if_fp4, quantize_mha): def filter_func_ltx_video(name: str) -> bool: """Filter function specifically for LTX-Video models.""" pattern = re.compile( - r".*(proj_in|time_embed|caption_projection|proj_out|patchify_proj|adaln_single|transformer_blocks\.(0|1|2|46|47)\.).*" + r".*(proj_in|time_embed|caption_projection|proj_out|patchify_proj|adaln_single|transformer_blocks\.(0|1|2|45|46|47)\.).*" ) return pattern.match(name) is not None