From df8a3fd5c1e8ff6b9d0cafc715deae496d6194fd Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Tue, 1 Sep 2026 17:28:06 +0800 Subject: [PATCH] Fix pre-weights ignoring dit_quant_scheme (fp8 dtype mismatch) ZImagePreWeights hard-coded its matmul layers (img_in / txt_in / timestep embedder) to the "Default" scheme, unlike ZImagePostWeights and ZImageTransformerWeights which follow config["dit_quant_scheme"]. With a pre-quantized FP8 checkpoint (e.g. qwen_image_edit_2509_fp8_e4m3fn_scaled.safetensors), all_x_embedder is stored as float8_e4m3fn but img_in still runs the Default addmm path, so the bf16 activation and the fp8 weight fail with: RuntimeError: self and mat2 must have the same dtype, but got BFloat16 and Float8_e4m3fn Align pre-weights with post/transformer weights: resolve mm_type once from config and register every matmul layer via MM_WEIGHT_REGISTER[self.mm_type]. Fixes #1439 --- lightx2v/models/networks/z_image/weights/pre_weights.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lightx2v/models/networks/z_image/weights/pre_weights.py b/lightx2v/models/networks/z_image/weights/pre_weights.py index 0e0299b52..cfda32fc5 100755 --- a/lightx2v/models/networks/z_image/weights/pre_weights.py +++ b/lightx2v/models/networks/z_image/weights/pre_weights.py @@ -10,18 +10,19 @@ class ZImagePreWeights(WeightModule): def __init__(self, config): super().__init__() self.config = config + self.mm_type = config.get("dit_quant_scheme", "Default") self.add_module( "img_in", - MM_WEIGHT_REGISTER["Default"]("all_x_embedder.2-1.weight", "all_x_embedder.2-1.bias"), + MM_WEIGHT_REGISTER[self.mm_type]("all_x_embedder.2-1.weight", "all_x_embedder.2-1.bias"), ) self.add_module( "txt_in", - MM_WEIGHT_REGISTER["Default"]("cap_embedder.1.weight", "cap_embedder.1.bias"), + MM_WEIGHT_REGISTER[self.mm_type]("cap_embedder.1.weight", "cap_embedder.1.bias"), ) self.add_module("txt_norm", RMS_WEIGHT_REGISTER["torch"]("cap_embedder.0.weight")) - self.add_module("time_text_embed_timestep_embedder_linear_1", MM_WEIGHT_REGISTER["Default"]("t_embedder.mlp.0.weight", "t_embedder.mlp.0.bias")) - self.add_module("time_text_embed_timestep_embedder_linear_2", MM_WEIGHT_REGISTER["Default"]("t_embedder.mlp.2.weight", "t_embedder.mlp.2.bias")) + self.add_module("time_text_embed_timestep_embedder_linear_1", MM_WEIGHT_REGISTER[self.mm_type]("t_embedder.mlp.0.weight", "t_embedder.mlp.0.bias")) + self.add_module("time_text_embed_timestep_embedder_linear_2", MM_WEIGHT_REGISTER[self.mm_type]("t_embedder.mlp.2.weight", "t_embedder.mlp.2.bias")) self.add_module("x_pad_token", TENSOR_REGISTER["Default"]("x_pad_token")) self.add_module("cap_pad_token", TENSOR_REGISTER["Default"]("cap_pad_token"))