receive any kinds of image extension in SanaImgDataset;

Signed-off-by: lawrence-cj <[email protected]>
NVlabs · Dec 17, 2024 · dae74bf · dae74bf
1 parent 6bfa5b8
commit dae74bf
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -5,7 +5,6 @@
 .gradio/
 .idea/
 *.png
-*.json
 tmp*
 output*
 output/

diff --git a/asset/example_data/00000000.jpg b/asset/example_data/00000000.jpg
diff --git a/asset/example_data/meta_data.json b/asset/example_data/meta_data.json
@@ -2,6 +2,6 @@
     "name": "sana-dev",
     "__kind__": "Sana-ImgDataset",
     "img_names": [
-        "00000000", "00000000"
+        "00000000", "00000000", "00000000.png", "00000000.jpg"
     ]
 }
diff --git a/diffusion/data/datasets/sana_data.py b/diffusion/data/datasets/sana_data.py
@@ -92,8 +92,12 @@ def __init__(
 
     def getdata(self, idx):
         data = self.dataset[idx]
+        img_extensions = [".jpg", ".png", ".jpeg", ".webp"]
+        filename, ext = os.path.splitext(data)
+        if ext in img_extensions:
+            data = filename
+            self.img_extension = ext
         self.key = data.split("/")[-1]
-        # info = json.load(open(f"{data}.json"))[self.key]
         info = {}
         with open(f"{data}.txt") as f:
             info[self.default_prompt] = f.readlines()[0].strip()
-Original file line number
+Diff line change
@@ Expand Up / @@ -5,7 +5,6 @@ @@
     .gradio/
     .idea/
     *.png
-    *.json
     tmp*
     output*
     output/
@@ Expand Down @@