From 1f87f83e6866351be76289ce7a8d0d8efc0636f0 Mon Sep 17 00:00:00 2001
From: aihao <51043929+aihao2000@users.noreply.github.com>
Date: Fri, 10 Nov 2023 17:03:14 +0800
Subject: [PATCH] add load_datasete data_dir parameter (#5747)

---
 examples/text_to_image/train_text_to_image_flax.py      | 1 +
 examples/text_to_image/train_text_to_image_lora_sdxl.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/examples/text_to_image/train_text_to_image_flax.py b/examples/text_to_image/train_text_to_image_flax.py
index 64b71b4f83..e5410a0a9e 100644
--- a/examples/text_to_image/train_text_to_image_flax.py
+++ b/examples/text_to_image/train_text_to_image_flax.py
@@ -275,6 +275,7 @@ def main():
             args.dataset_name,
             args.dataset_config_name,
             cache_dir=args.cache_dir,
+            data_dir=args.train_data_dir
         )
     else:
         data_files = {}
diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py
index f0d83d55e9..d621858e3d 100644
--- a/examples/text_to_image/train_text_to_image_lora_sdxl.py
+++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py
@@ -768,6 +768,7 @@ def main(args):
             args.dataset_name,
             args.dataset_config_name,
             cache_dir=args.cache_dir,
+            data_dir=args.train_data_dir
         )
     else:
         data_files = {}