Support both str and Path jsonl_path inputs in ImageCaptionJsonlDataset.

invoke-ai · Feb 12, 2024 · df50169 · df50169
1 parent d2c185e
commit df50169
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py b/src/invoke_training/_shared/data/datasets/image_caption_jsonl_dataset.py
@@ -21,10 +21,13 @@ class ImageCaptionJsonlDataset(torch.utils.data.Dataset):
     """A dataset that loads images and captions from a directory of image files and .txt files."""
 
     def __init__(
-        self, jsonl_path: Path, image_column: str = IMAGE_COLUMN_DEFAULT, caption_column: str = CAPTION_COLUMN_DEFAULT
+        self,
+        jsonl_path: Path | str,
+        image_column: str = IMAGE_COLUMN_DEFAULT,
+        caption_column: str = CAPTION_COLUMN_DEFAULT,
     ):
         super().__init__()
-        self._jsonl_path = jsonl_path
+        self._jsonl_path = Path(jsonl_path)
         self._image_column = image_column
         self._caption_column = caption_column