From 40b7d2bae4ee2c37eecee6565a383abed20e9600 Mon Sep 17 00:00:00 2001
From: yunxiang wei <121092793+yiwangsimple@users.noreply.github.com>
Date: Wed, 17 Jul 2024 23:52:13 +0800
Subject: [PATCH] Add files via upload

---
 SD3LongCaptioner.py    | 3 +++
 SD3LongCaptioner_v2.py | 5 +++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/SD3LongCaptioner.py b/SD3LongCaptioner.py
index 96fda0d..c4f14a3 100644
--- a/SD3LongCaptioner.py
+++ b/SD3LongCaptioner.py
@@ -63,6 +63,9 @@ class SD3LongCaptioner:
         # 解码生成的文本
         decoded = self.processor.decode(generation[0], skip_special_tokens=True)
         
+        # 移除开头的提示词（如果存在）
+        decoded = decoded.replace(prompt, "", 1).strip()
+        
         return (decoded,)
 
 NODE_CLASS_MAPPINGS = {
diff --git a/SD3LongCaptioner_v2.py b/SD3LongCaptioner_v2.py
index f9fe859..8e64b8d 100644
--- a/SD3LongCaptioner_v2.py
+++ b/SD3LongCaptioner_v2.py
@@ -11,8 +11,6 @@ from huggingface_hub import snapshot_download
 files_for_sd3_long_captioner_v2 = Path(os.path.join(folder_paths.models_dir, "LLavacheckpoints", "files_for_sd3_long_captioner_v2"))
 files_for_sd3_long_captioner_v2.mkdir(parents=True, exist_ok=True)
 
-# ... 其余代码 ...
-
 class SD3LongCaptionerV2:
     def __init__(self):
         self.model_id = "gokaygokay/sd3-long-captioner-v2"
@@ -65,6 +63,9 @@ class SD3LongCaptionerV2:
         # 解码生成的文本
         decoded = self.processor.decode(generation[0], skip_special_tokens=True)
         
+        # 移除开头的提示词（如果存在）
+        decoded = decoded.replace(prompt, "", 1).strip()
+        
         return (decoded,)
 
 NODE_CLASS_MAPPINGS = {