Adds common natural img and remove previous

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.pth.tar filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+*.JPEG filter=lfs diff=lfs merge=lfs -text

datasets.py

@@ -93,7 +93,7 @@ class LsdirMiniDataset(torch.utils.data.Dataset):
 	transform: Optional[Callable] = None,
     ) -> None:
         self.root = root
-        self.image_files = [f for f in os.listdir(self.root) if f.lower().endswith('.png')]
+        self.image_files = [f for f in os.listdir(self.root) if f.lower().endswith(('.png', '.JPEG'))]
         self.transform = transform
 
     def __len__(self) -> int:
@@ -101,7 +101,7 @@ class LsdirMiniDataset(torch.utils.data.Dataset):
 
     def __getitem__(self, idx):
         img_path = os.path.join(self.root, self.image_files[idx])
-        img = Image.open(img_path)
+        img = Image.open(img_path).convert("RGB")  # Ensure consistent 3-channel format
         if self.transform:
             img = self.transform(img)