From 5ea5fefcd2a5d67e885d8fce23c26911171782ad Mon Sep 17 00:00:00 2001
From: Yuta Hayashibe <yuta@hayashibe.jp>
Date: Mon, 6 Feb 2023 12:29:41 +0900
Subject: [PATCH] Always join with ", "

---
 library/train_util.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index bb463098..278d2f8d 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -135,7 +135,7 @@ class BaseDataset(torch.utils.data.Dataset):
 
   def process_caption(self, caption):
     if self.shuffle_caption:
-      tokens = caption.strip().split(",")
+      tokens = [t.strip() for t in caption.strip().split(",")]
       if self.shuffle_keep_tokens is None:
         random.shuffle(tokens)
       else:
@@ -144,7 +144,7 @@ class BaseDataset(torch.utils.data.Dataset):
           tokens = tokens[self.shuffle_keep_tokens:]
           random.shuffle(tokens)
           tokens = keep_tokens + tokens
-      caption = ",".join(tokens).strip()
+      caption = ", ".join(tokens)
 
     for str_from, str_to in self.replacements.items():
       if str_from == "":