From a8ea5a6ac4b680b14a0bcffc6eaf32ec3f9c6442 Mon Sep 17 00:00:00 2001
From: DKnight54 <126916963+DKnight54@users.noreply.github.com>
Date: Fri, 31 Jan 2025 03:25:34 +0800
Subject: [PATCH] Update train_util.py

---
 library/train_util.py | 33 ++++++++++++---------------------
 1 file changed, 12 insertions(+), 21 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index 83ae5122..4d644f9e 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -5507,10 +5507,6 @@ def sample_images_common(
 
     # preprocess prompts
     idx = 0
-    if example_tuple:
-        logger.info(f"len(example_tuple): {len(example_tuple)}")
-        logger.info(f"len(example_tuple[0]): {len(example_tuple[0])}")
-        logger.info(f"len(example_tuple[1]): {len(example_tuple[1])}")
     for i in range(len(prompts)):
         prompt_dict = prompts[i]
         if isinstance(prompt_dict, str):
@@ -5520,23 +5516,18 @@ def sample_images_common(
         
         if '__caption__' in prompt_dict.get("prompt") and example_tuple:
             logger.info(f"Original prompt: {prompt_dict.get('prompt')}")
-            if len(example_tuple[1]) > 1:
-                while example_tuple[1][idx] == '':
-                    idx = (idx + 1) % len(example_tuple[1])
-                    if idx == 0:
-                        break
-                    prompt_dict["prompt"] = f"{example_tuple[1][idx]}"
-                    logger.info(f"Replacement prompt: {example_tuple[1][idx]}")
-                    prompt_dict["height"] = example_tuple[0].shape[2] * 8
-                    prompt_dict["width"] = example_tuple[0].shape[3] * 8
-                    prompt_dict["original_lantent"] = example_tuple[0][idx].unsqueeze(0)
-                    idx = (idx + 1) % len(example_tuple[1])
-            else:
-                prompt_dict["prompt"] = f"{example_tuple[1]}"
-                logger.info(f"Replacement prompt: {example_tuple[1]}")
-                prompt_dict["height"] = example_tuple[0].shape[2] * 8
-                prompt_dict["width"] = example_tuple[0].shape[3] * 8
-                prompt_dict["original_lantent"] = example_tuple[0].unsqueeze(0)
+
+            while example_tuple[1][idx] == '':
+                idx = (idx + 1) % len(example_tuple[1])
+                if idx == 0:
+                    break
+            prompt_dict["prompt"] = f"{example_tuple[1][idx]}"
+            logger.info(f"Replacement prompt: {example_tuple[1][idx]}")
+            prompt_dict["height"] = example_tuple[0].shape[2] * 8
+            prompt_dict["width"] = example_tuple[0].shape[3] * 8
+            prompt_dict["original_lantent"] = example_tuple[0][idx].unsqueeze(0)
+            idx = (idx + 1) % len(example_tuple[1])
+
         prompt_dict["enum"] = i
         prompt_dict.pop("subset", None)
         prompts[i] = prompt_dict