fix: prompt unaligned loading

jolibrain · Jun 20, 2024 · e25d4b1 · e25d4b1
1 parent 2538c00
commit e25d4b1
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/data/unaligned_dataset.py b/data/unaligned_dataset.py
@@ -62,9 +62,12 @@ def get_img(
         A = self.transform_A(A_img)
         B = self.transform_B(B_img)
 
-        real_B_prompt = self.B_img_prompt[B_img_path]
-        if len(real_B_prompt) == 1 and isinstance(real_B_prompt[0], str):
-            real_B_prompt = real_B_prompt[0]
+        if self.B_img_prompt is not None:
+            real_B_prompt = self.B_img_prompt[B_img_path]
+            if len(real_B_prompt) == 1 and isinstance(real_B_prompt[0], str):
+                real_B_prompt = real_B_prompt[0]
+        else:
+            real_B_prompt = None
 
         return {
             "A": A,