add option to use batch size for training

1 files changed, 3 insertions, 0 deletions

diff --git a/modules/ui.py b/modules/ui.py
index 1bc919c7..45550ea8 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1166,6 +1166,7 @@ def create_ui(wrap_gradio_gpu_call):
                     train_embedding_name = gr.Dropdown(label='Embedding', choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()))

                     train_hypernetwork_name = gr.Dropdown(label='Hypernetwork', choices=[x for x in shared.hypernetworks.keys()])

                     learn_rate = gr.Textbox(label='Learning rate', placeholder="Learning rate", value="0.005")

+                    batch_size = gr.Number(label='Batch size', value=1, precision=0)

                     dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")

                     log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")

                     template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))

@@ -1244,6 +1245,7 @@ def create_ui(wrap_gradio_gpu_call):
             inputs=[

                 train_embedding_name,

                 learn_rate,

+                batch_size,

                 dataset_directory,

                 log_directory,

                 training_width,

@@ -1268,6 +1270,7 @@ def create_ui(wrap_gradio_gpu_call):
             inputs=[

                 train_hypernetwork_name,

                 learn_rate,

+                batch_size,

                 dataset_directory,

                 log_directory,

                 steps,