Update caption extension fields to Dropdown with restricted choices f…

…or security reasons
bmaltais · Apr 11, 2024 · bc7403d · bc7403d
1 parent fc29524
commit bc7403d
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 18 deletions.
diff --git a/kohya_gui/basic_caption_gui.py b/kohya_gui/basic_caption_gui.py
@@ -190,9 +190,9 @@ def list_images_dirs(path):
                 show_progress=False,
             )
             # Textbox for caption file extension
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )

diff --git a/kohya_gui/blip_caption_gui.py b/kohya_gui/blip_caption_gui.py
@@ -142,9 +142,9 @@ def list_train_dirs(path):
                 show_progress=False,
             )
         with gr.Row():
-            caption_file_ext = gr.Textbox(
+            caption_file_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )

diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py
@@ -98,10 +98,11 @@ def init_training_controls(self) -> None:
                 label="Save every N epochs", value=self.config.get("basic.save_every_n_epochs", 1), precision=0
             )
             # Initialize the caption extension input
-            self.caption_extension = gr.Textbox(
-                label="Caption Extension",
-                placeholder="(Optional) default: .caption",
-                value=self.config.get("basic.caption_extension", ""),
+            self.caption_extension = gr.Dropdown(
+                label="Caption file extension",
+                choices=[".cap", ".caption", ".txt"],
+                value=".txt",
+                interactive=True,
             )
 
     def init_precision_and_resources_controls(self) -> None:

diff --git a/kohya_gui/git_caption_gui.py b/kohya_gui/git_caption_gui.py
@@ -114,9 +114,9 @@ def list_train_dirs(path):
                 show_progress=False,
             )
         with gr.Row():
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )

diff --git a/kohya_gui/group_images_gui.py b/kohya_gui/group_images_gui.py
@@ -171,9 +171,9 @@ def list_output_dirs(path):
                 info="Generate caption files for the grouped images based on their folder name",
             )
 
-            caption_ext = gr.Textbox(
-                label="Caption Extension",
-                placeholder="Caption file extension (e.g., .txt)",
+            caption_ext = gr.Dropdown(
+                label="Caption file extension",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )

diff --git a/kohya_gui/manual_caption_gui.py b/kohya_gui/manual_caption_gui.py
@@ -298,9 +298,9 @@ def list_images_dirs(path):
                 show_progress=False,
             )
             load_images_button = gr.Button("Load", elem_id="open_folder")
-            caption_ext = gr.Textbox(
+            caption_ext = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
+                choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
             )

diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py
@@ -188,10 +188,10 @@ def list_train_dirs(path):
 
         with gr.Row():
 
-            caption_extension = gr.Textbox(
+            caption_extension = gr.Dropdown(
                 label="Caption file extension",
-                placeholder="Extension for caption file (e.g., .caption, .txt)",
-                value=config.get("wd14_caption.caption_extension", ".txt"),
+                choices=[".cap", ".caption", ".txt"],
+                value=".txt",
                 interactive=True,
             )