From a7f31c6199cd12fa3104cc2898335782728cd7d0 Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Sun, 5 Jul 2026 15:29:38 +0800
Subject: [PATCH] fix: avoid layer splitting unet block paths

---
 src/core/layer_split_partition.cpp | 47 +++++++++++++++++++++++-------
 1 file changed, 37 insertions(+), 10 deletions(-)

diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp
index 8e33eb98a..951965e18 100644
--- a/src/core/layer_split_partition.cpp
+++ b/src/core/layer_split_partition.cpp
@@ -9,21 +9,48 @@
 
 namespace sd {
 
+    static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
+        return pos == 0 || name[pos - 1] == '.';
+    }
+
+    static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
+        size_t pos = name.find(segment);
+        while (pos != std::string::npos) {
+            if (layer_split_path_segment_starts_at(name, pos)) {
+                return true;
+            }
+            pos = name.find(segment, pos + 1);
+        }
+        return false;
+    }
+
     int layer_split_tensor_block_index(const std::string& name) {
+        static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
+                                                    "down_blocks.", "up_blocks.", "mid_block."};
+        for (const char* segment : unet_block_segments) {
+            if (layer_split_has_path_segment(name, segment)) {
+                return -1;
+            }
+        }
+
         static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
                                                "single_blocks.", "blocks.", "block.", "layers."};
         for (const char* keyword : block_keywords) {
             size_t pos = name.find(keyword);
-            if (pos == std::string::npos) {
-                continue;
-            }
-            pos += std::strlen(keyword);
-            size_t end = pos;
-            while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
-                end++;
-            }
-            if (end > pos && (end == name.size() || name[end] == '.')) {
-                return std::atoi(name.substr(pos, end - pos).c_str());
+            while (pos != std::string::npos) {
+                if (!layer_split_path_segment_starts_at(name, pos)) {
+                    pos = name.find(keyword, pos + 1);
+                    continue;
+                }
+                pos += std::strlen(keyword);
+                size_t end = pos;
+                while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
+                    end++;
+                }
+                if (end > pos && (end == name.size() || name[end] == '.')) {
+                    return std::atoi(name.substr(pos, end - pos).c_str());
+                }
+                break;
             }
         }
         return -1;