From 5bccd60c803e421d5d9ed6b6a449df4a4f17d9e2 Mon Sep 17 00:00:00 2001
From: Zhongjie Duan <35051019+Artiprocher@users.noreply.github.com>
Date: Mon, 23 Mar 2026 11:24:49 +0800
Subject: [PATCH] compatibility patch (#1363)

---
 diffsynth/configs/model_configs.py                   | 7 +++++++
 diffsynth/utils/state_dict_converters/z_image_dit.py | 3 +++
 2 files changed, 10 insertions(+)
 create mode 100644 diffsynth/utils/state_dict_converters/z_image_dit.py

diff --git a/diffsynth/configs/model_configs.py b/diffsynth/configs/model_configs.py
index 9593f0b..4222202 100644
--- a/diffsynth/configs/model_configs.py
+++ b/diffsynth/configs/model_configs.py
@@ -604,6 +604,13 @@ z_image_series = [
         "extra_kwargs": {"model_size": "0.6B"},
         "state_dict_converter": "diffsynth.utils.state_dict_converters.z_image_text_encoder.ZImageTextEncoderStateDictConverter",
     },
+    {
+        # To ensure compatibility with the `model.diffusion_model` prefix introduced by other frameworks.
+        "model_hash": "8cf241a0d32f93d5de368502a086852f",
+        "model_name": "z_image_dit",
+        "model_class": "diffsynth.models.z_image_dit.ZImageDiT",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.z_image_dit.ZImageDiTStateDictConverter",
+    },
 ]
 """
 Offical model repo: https://www.modelscope.cn/models/Lightricks/LTX-2
diff --git a/diffsynth/utils/state_dict_converters/z_image_dit.py b/diffsynth/utils/state_dict_converters/z_image_dit.py
new file mode 100644
index 0000000..0f44d8b
--- /dev/null
+++ b/diffsynth/utils/state_dict_converters/z_image_dit.py
@@ -0,0 +1,3 @@
+def ZImageDiTStateDictConverter(state_dict):
+    state_dict_ = {name.replace("model.diffusion_model.", ""): state_dict[name] for name in state_dict}
+    return state_dict_