From b0b1337bc5e5ca66fd2e521b1dd671e70c238a63 Mon Sep 17 00:00:00 2001
From: cafeai <116491182+cafeai@users.noreply.github.com>
Date: Sun, 6 Nov 2022 00:32:15 +0900
Subject: [PATCH] ImageStore Filename Fix

---
 diffusers_trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/diffusers_trainer.py b/diffusers_trainer.py
index 2171a8d..3a59c5a 100644
--- a/diffusers_trainer.py
+++ b/diffusers_trainer.py
@@ -20,6 +20,7 @@ import time
 import itertools
 import numpy as np
 import json
+import re
 
 try:
     pynvml.nvmlInit()
@@ -171,7 +172,7 @@ class ImageStore:
     
     # gets caption by removing the extension from the filename and replacing it with .txt
     def get_caption(self, index: int) -> str:
-        filename = self.image_files[index].split('.')[0] + '.txt'
+        filename = re.sub('\.[^/.]+$', '', self.image_files[index]) + '.txt'
         with open(filename, 'r') as f:
             return f.read()