From 1aab32fe034ed98fdd2ba2f8b18eccb48f7a2f63 Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Tue, 9 Dec 2025 18:09:47 +0800
Subject: [PATCH] fixed safetensors loading for zimage

---
 otherarch/sdcpp/stable-diffusion.cpp | 35 ++++++++++++++--------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/otherarch/sdcpp/stable-diffusion.cpp b/otherarch/sdcpp/stable-diffusion.cpp
index f691589cf..10ce0c8b7 100644
--- a/otherarch/sdcpp/stable-diffusion.cpp
+++ b/otherarch/sdcpp/stable-diffusion.cpp
@@ -264,6 +264,24 @@ public:
 
         bool is_unet = sd_version_is_unet(model_loader.get_sd_version());
         int tempver = model_loader.get_sd_version();
+
+        // kcpp fallback to separate diffusion model passed as model
+        if (tempver == VERSION_COUNT &&
+        strlen(SAFE_STR(sd_ctx_params->model_path)) > 0 &&
+        strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) == 0 &&
+        (t5_path_fixed!=""||clipl_path_fixed!=""))
+        {
+            bool endswithsafetensors = ends_with(sd_ctx_params->model_path, ".safetensors");
+            if(endswithsafetensors && !model_loader.has_diffusion_model_tensors())
+            {
+                LOG_INFO("SD Diffusion Model tensors missing! Fallback trying alternative tensor names...\n");
+                if (!model_loader.init_from_file(sd_ctx_params->model_path, "model.diffusion_model.")) {
+                    LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->model_path);
+                }
+                tempver = model_loader.get_sd_version();
+            }
+        }
+
         bool iswan = (tempver==VERSION_WAN2 || tempver==VERSION_WAN2_2_I2V || tempver==VERSION_WAN2_2_TI2V);
         bool isqwenimg = (tempver==VERSION_QWEN_IMAGE);
         bool iszimg = (tempver==VERSION_Z_IMAGE);
@@ -370,23 +388,6 @@ public:
 
         version = model_loader.get_sd_version();
 
-        // kcpp fallback to separate diffusion model passed as model
-        if (version == VERSION_COUNT &&
-            strlen(SAFE_STR(sd_ctx_params->model_path)) > 0 &&
-            strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) == 0 &&
-            t5_path_fixed!="" )
-        {
-            bool endswithsafetensors = ends_with(sd_ctx_params->model_path, ".safetensors");
-            if(endswithsafetensors && !model_loader.has_diffusion_model_tensors())
-            {
-                LOG_INFO("SD Diffusion Model tensors missing! Fallback trying alternative tensor names...\n");
-                if (!model_loader.init_from_file(sd_ctx_params->model_path, "model.diffusion_model.")) {
-                    LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->model_path);
-                }
-                version = model_loader.get_sd_version();
-            }
-        }
-
         if (version == VERSION_COUNT) {
             LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path));
             return false;