From 60b374b904eaaf37f1181192b92a2515eb3d1ced Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Wed, 19 Jun 2024 03:05:35 +0800
Subject: [PATCH] 'update'

---
 mz_train_tools_core.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py
index f7c1975..9625000 100644
--- a/mz_train_tools_core.py
+++ b/mz_train_tools_core.py
@@ -30,12 +30,14 @@ def MZ_KohyaSSInitWorkspace_call(args={}):
         branch = args.get("branch", "main")
 
         # 查看本地分支是否一致
+        short_result = subprocess.run(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
         result = subprocess.run(
-            ["git", "branch"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
+            ["git", "rev-parse", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True)
+        print(
+            f"current branch: {short_result.stdout.decode()} {result.stdout.decode()}")
 
-        print(f"current branch: {result.stdout.decode()}")
-
-        if branch not in result.stdout.decode():
+        if branch != result.stdout.decode() and branch != short_result.stdout.decode():
             subprocess.run(
                 ["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
             subprocess.run(