Update waifuc_gui/interface.py

waifuc_gui/interface.py

@@ -1,244 +1,251 @@
-import gradio as gr
-from .source_manager import SourceManager
-from .action_manager import ActionManager
-from .exporter_manager import ExporterManager
-from .config_manager import ConfigManager
-
-class Interface:
-    def __init__(self, source_manager: SourceManager, action_manager: ActionManager, exporter_manager: ExporterManager, config_manager: ConfigManager):
-        self.source_manager = source_manager
-        self.action_manager = action_manager
-        self.exporter_manager = exporter_manager
-        self.config_manager = config_manager
-        self.params = source_manager.create_param_inputs()
-        self.param_components = [param_input for source_params in self.params.values() for param_input in source_params.values()]
-        self.language_dict = {
-            "zh": {
-                "title": "Waifuc 数据收集工具",
-                "data_collection": "数据收集",
-                "config": "配置",
-                "logs": "日志",
-                "select_source": "选择数据源",
-                "select_actions": "选择动作",
-                "dataset_name": "数据集名称",
-                "select_exporter": "选择导出器",
-                "start_collection": "开始收集",
-                "status": "状态",
-                "download_data": "下载收集的数据",
-                "config_management": "配置管理",
-                "pixiv_token": "Pixiv刷新令牌",
-                "output_dir": "输出目录",
-                "num_items": "默认图片数量",
-                "resize_size": "默认调整大小",
-                "language": "语言",
-                "save_config": "保存配置",
-                "export_config": "导出配置",
-                "import_config": "导入配置",
-                "config_status": "配置状态",
-                "view_logs": "查看日志",
-                "download_log": "下载日志"
-            },
-            "en": {
-                "title": "Waifuc Data Collection Tool",
-                "data_collection": "Data Collection",
-                "config": "Configuration",
-                "logs": "Logs",
-                "select_source": "Select Data Source",
-                "select_actions": "Select Actions",
-                "dataset_name": "Dataset Name",
-                "select_exporter": "Select Exporter",
-                "start_collection": "Start Collection",
-                "status": "Status",
-                "download_data": "Download Collected Data",
-                "config_management": "Configuration Management",
-                "pixiv_token": "Pixiv Refresh Token",
-                "output_dir": "Output Directory",
-                "num_items": "Default Number of Items",
-                "resize_size": "Default Resize Size",
-                "language": "Language",
-                "save_config": "Save Configuration",
-                "export_config": "Export Configuration",
-                "import_config": "Import Configuration",
-                "config_status": "Configuration Status",
-                "view_logs": "View Logs",
-                "download_log": "Download Log"
-            }
-        }
-
-    def get_text(self, key):
-        return self.language_dict[self.config_manager.get_config("language")][key]
-
-    def update_params_visibility(self, selected_source):
-        updates = []
-        for source, source_params in self.params.items():
-            for param_name, param_input in source_params.items():
-                updates.append(param_input.update(visible=(source == selected_source)))
-        return updates
-
-    def collect_params(self, selected_source, *param_values):
-        collected = {}
-        source_params = self.params.get(selected_source, {})
-        for param_name, param_value in zip(source_params.keys(), param_values):
-            collected[param_name] = param_value
-        return collected
-
-    def update_action_params(self, selected_actions):
-        action_params = self.action_manager.create_action_param_inputs(selected_actions)
-        components = []
-        for action, params in action_params.items():
-            for param_name, param_input in params.items():
-                components.append(param_input)
-        return components, action_params
-
-    def build(self):
-        with gr.Blocks(title=self.get_text("title")) as demo:
-            language_dropdown = gr.Dropdown(
-                choices=["zh", "en"],
-                label=self.get_text("language"),
-                value=self.config_manager.get_config("language")
-            )
-            with gr.Tab(self.get_text("data_collection")):
-                gr.Markdown("### " + self.get_text("data_collection"))
-                source_dropdown = gr.Dropdown(
-                    choices=self.source_manager.source_names,
-                    label=self.get_text("select_source"),
-                    value=self.source_manager.source_names[0] if self.source_manager.source_names else None,
-                    info="Choose a data source like DanbooruSource or LocalSource" if self.config_manager.get_config("language") == "en" else "选择数据源，如DanbooruSource或LocalSource"
-                )
-                param_components = self.param_components
-                action_checkboxes = gr.CheckboxGroup(
-                    choices=self.action_manager.action_names,
-                    label=self.get_text("select_actions"),
-                    value=[],
-                    info="Select actions to process images, e.g., NoMonochromeAction" if self.config_manager.get_config("language") == "en" else "选择处理图片的动作，例如NoMonochromeAction"
-                )
-                action_param_components = gr.State(value={})
-                action_param_inputs = []
-                dataset_name_input = gr.Textbox(
-                    label=self.get_text("dataset_name"),
-                    value="waifuc_dataset",
-                    info="Name for the output dataset" if self.config_manager.get_config("language") == "en" else "输出数据集的名称"
-                )
-                exporter_dropdown = gr.Dropdown(
-                    choices=self.exporter_manager.exporter_names,
-                    label=self.get_text("select_exporter"),
-                    value=self.exporter_manager.exporter_names[0] if self.exporter_manager.exporter_names else None,
-                    info="Choose an exporter like TextualInversionExporter" if self.config_manager.get_config("language") == "en" else "选择导出器，如TextualInversionExporter"
-                )
-                start_btn = gr.Button(self.get_text("start_collection"))
-                status = gr.Textbox(label=self.get_text("status"), interactive=False)
-                output_file = gr.File(label=self.get_text("download_data"))
-            with gr.Tab(self.get_text("config")):
-                gr.Markdown("### " + self.get_text("config_management"))
-                pixiv_token_input = gr.Textbox(
-                    label=self.get_text("pixiv_token"),
-                    value=self.config_manager.get_config("pixiv_refresh_token"),
-                    type="password",
-                    info="Required for PixivSearchSource, stored in session only" if self.config_manager.get_config("language") == "en" else "PixivSearchSource所需，仅存储在会话中"
-                )
-                output_dir_input = gr.Textbox(
-                    label=self.get_text("output_dir"),
-                    value=self.config_manager.get_config("output_dir"),
-                    info="Directory for output data, session-specific" if self.config_manager.get_config("language") == "en" else "输出数据的目录，会话专用"
-                )
-                num_items_input = gr.Number(
-                    label=self.get_text("num_items"),
-                    value=self.config_manager.get_config("default_num_items"),
-                    info="Default number of images to collect" if self.config_manager.get_config("language") == "en" else "默认收集的图片数量"
-                )
-                resize_size_input = gr.Number(
-                    label=self.get_text("resize_size"),
-                    value=self.config_manager.get_config("default_resize_size"),
-                    info="Default size for ResizeAction" if self.config_manager.get_config("language") == "en" else "ResizeAction的默认尺寸"
-                )
-                config_export_btn = gr.Button(self.get_text("export_config"))
-                config_import_file = gr.File(label=self.get_text("import_config"))
-                save_config_btn = gr.Button(self.get_text("save_config"))
-                config_status = gr.Textbox(label=self.get_text("config_status"), interactive=False)
-            with gr.Tab(self.get_text("logs")):
-                gr.Markdown("### " + self.get_text("view_logs"))
-                log_output = gr.Textbox(label=self.get_text("view_logs"), interactive=False, lines=10)
-                log_download_btn = gr.Button(self.get_text("download_log"))
-                log_download_file = gr.File(label=self.get_text("download_log"))
-            def update_language(language):
-                self.config_manager.set_config("language", language)
-                return {
-                    "demo": demo.update(title=self.get_text("title")),
-                    "language_dropdown": language_dropdown.update(label=self.get_text("language")),
-                    "source_dropdown": source_dropdown.update(label=self.get_text("select_source")),
-                    "action_checkboxes": action_checkboxes.update(label=self.get_text("select_actions")),
-                    "dataset_name_input": dataset_name_input.update(label=self.get_text("dataset_name")),
-                    "exporter_dropdown": exporter_dropdown.update(label=self.get_text("select_exporter")),
-                    "start_btn": start_btn.update(value=self.get_text("start_collection")),
-                    "status": status.update(label=self.get_text("status")),
-                    "output_file": output_file.update(label=self.get_text("download_data")),
-                    "pixiv_token_input": pixiv_token_input.update(label=self.get_text("pixiv_token")),
-                    "output_dir_input": output_dir_input.update(label=self.get_text("output_dir")),
-                    "num_items_input": num_items_input.update(label=self.get_text("num_items")),
-                    "resize_size_input": resize_size_input.update(label=self.get_text("resize_size")),
-                    "config_export_btn": config_export_btn.update(value=self.get_text("export_config")),
-                    "config_import_file": config_import_file.update(label=self.get_text("import_config")),
-                    "save_config_btn": save_config_btn.update(value=self.get_text("save_config")),
-                    "config_status": config_status.update(label=self.get_text("config_status")),
-                    "log_output": log_output.update(label=self.get_text("view_logs")),
-                    "log_download_btn": log_download_btn.update(value=self.get_text("download_log")),
-                    "log_download_file": log_download_file.update(label=self.get_text("download_log"))
-                }
-            language_dropdown.change(
-                fn=update_language,
-                inputs=language_dropdown,
-                outputs=[
-                    demo, language_dropdown, source_dropdown, action_checkboxes,
-                    dataset_name_input, exporter_dropdown, start_btn, status,
-                    output_file, pixiv_token_input, output_dir_input, num_items_input,
-                    resize_size_input, config_export_btn, config_import_file,
-                    save_config_btn, config_status, log_output, log_download_btn,
-                    log_download_file
-                ]
-            )
-            source_dropdown.change(
-                fn=self.update_params_visibility,
-                inputs=source_dropdown,
-                outputs=param_components
-            )
-            action_checkboxes.change(
-                fn=self.update_action_params,
-                inputs=action_checkboxes,
-                outputs=[gr.State(), gr.State()]
-            )
-            def save_configs(pixiv_token, output_dir, num_items, resize_size):
-                self.config_manager.set_config("pixiv_refresh_token", pixiv_token)
-                self.config_manager.set_config("output_dir", output_dir)
-                self.config_manager.set_config("default_num_items", int(num_items))
-                self.config_manager.set_config("default_resize_size", int(resize_size))
-                return self.get_text("config_status") + ": " + ("Configuration saved" if self.config_manager.get_config("language") == "en" else "配置已保存")
-            save_config_btn.click(
-                fn=save_configs,
-                inputs=[pixiv_token_input, output_dir_input, num_items_input, resize_size_input],
-                outputs=config_status
-            )
-            config_export_btn.click(
-                fn=self.config_manager.export_config,
-                inputs=[],
-                outputs=gr.File(label=self.get_text("export_config"))
-            )
-            config_import_file.upload(
-                fn=self.config_manager.import_config,
-                inputs=config_import_file,
-                outputs=config_status
-            )
-            return {
-                "demo": demo,
-                "source_dropdown": source_dropdown,
-                "param_components": param_components,
-                "action_checkboxes": action_checkboxes,
-                "action_param_components": action_param_components,
-                "dataset_name_input": dataset_name_input,
-                "exporter_dropdown": exporter_dropdown,
-                "start_btn": start_btn,
-                "status": status,
-                "output_file": output_file,
-                "log_output": log_output,
-                "log_download_btn": log_download_btn,
-                "log_download_file": log_download_file
-            }
+import gradio as gr
+from .source_manager import SourceManager
+from .action_manager import ActionManager
+from .exporter_manager import ExporterManager
+from .config_manager import ConfigManager
+
+class Interface:
+    def __init__(self, source_manager: SourceManager, action_manager: ActionManager, exporter_manager: ExporterManager, config_manager: ConfigManager):
+        self.source_manager = source_manager
+        self.action_manager = action_manager
+        self.exporter_manager = exporter_manager
+        self.config_manager = config_manager
+        self.params = source_manager.create_param_inputs()
+        self.param_components = [param_input for source_params in self.params.values() for param_input in source_params.values()]
+        self.language_dict = {
+            "zh": {
+                "title": "Waifuc 数据收集工具",
+                "data_collection": "数据收集",
+                "config": "配置",
+                "logs": "日志",
+                "select_source": "选择数据源",
+                "select_actions": "选择动作",
+                "dataset_name": "数据集名称",
+                "select_exporter": "选择导出器",
+                "start_collection": "开始收集",
+                "status": "状态",
+                "download_data": "下载收集的数据",
+                "config_management": "配置管理",
+                "pixiv_token": "Pixiv刷新令牌",
+                "output_dir": "输出目录",
+                "num_items": "默认图片数量",
+                "resize_size": "默认调整大小",
+                "language": "语言",
+                "save_config": "保存配置",
+                "export_config": "导出配置",
+                "import_config": "导入配置",
+                "config_status": "配置状态",
+                "view_logs": "查看日志",
+                "download_log": "下载日志",
+                "local_source_info": "上传包含图片的zip文件用于LocalSource"
+            },
+            "en": {
+                "title": "Waifuc Data Collection Tool",
+                "data_collection": "Data Collection",
+                "config": "Configuration",
+                "logs": "Logs",
+                "select_source": "Select Data Source",
+                "select_actions": "Select Actions",
+                "dataset_name": "Dataset Name",
+                "select_exporter": "Select Exporter",
+                "start_collection": "Start Collection",
+                "status": "Status",
+                "download_data": "Download Collected Data",
+                "config_management": "Configuration Management",
+                "pixiv_token": "Pixiv Refresh Token",
+                "output_dir": "Output Directory",
+                "num_items": "Default Number of Items",
+                "resize_size": "Default Resize Size",
+                "language": "Language",
+                "save_config": "Save Configuration",
+                "export_config": "Export Configuration",
+                "import_config": "Import Configuration",
+                "config_status": "Configuration Status",
+                "view_logs": "View Logs",
+                "download_log": "Download Log",
+                "local_source_info": "Upload a ZIP file containing images for LocalSource"
+            }
+        }
+
+    def get_text(self, key):
+        return self.language_dict[self.config_manager.get_config("language")][key]
+
+    def update_params_visibility(self, selected_source):
+        updates = []
+        for source, source_params in self.params.items():
+            for param_name, param_input in source_params.items():
+                updates.append(param_input.update(visible=(source == selected_source)))
+        local_source_info_visible = selected_source == "LocalSource"
+        return updates + [gr.Markdown.update(visible=local_source_info_visible)]
+
+    def collect_params(self, selected_source, *param_values):
+        collected = {}
+        source_params = self.params.get(selected_source, {})
+        for param_name, param_value in zip(source_params.keys(), param_values):
+            collected[param_name] = param_value
+        return collected
+
+    def update_action_params(self, selected_actions):
+        action_params = self.action_manager.create_action_param_inputs(selected_actions)
+        components = []
+        for action, params in action_params.items():
+            for param_name, param_input in params.items():
+                components.append(param_input)
+        return components, action_params
+
+    def build(self):
+        with gr.Blocks(title=self.get_text("title")) as demo:
+            language_dropdown = gr.Dropdown(
+                choices=["zh", "en"],
+                label=self.get_text("language"),
+                value=self.config_manager.get_config("language")
+            )
+            with gr.Tab(self.get_text("data_collection")):
+                gr.Markdown("### " + self.get_text("data_collection"))
+                source_dropdown = gr.Dropdown(
+                    choices=self.source_manager.source_names,
+                    label=self.get_text("select_source"),
+                    value=self.source_manager.source_names[0] if self.source_manager.source_names else None,
+                    info="Choose a data source like DanbooruSource or LocalSource" if self.config_manager.get_config("language") == "en" else "选择数据源，如DanbooruSource或LocalSource"
+                )
+                local_source_info = gr.Markdown(
+                    self.get_text("local_source_info"),
+                    visible=False
+                )
+                param_components = self.param_components
+                action_checkboxes = gr.CheckboxGroup(
+                    choices=self.action_manager.action_names,
+                    label=self.get_text("select_actions"),
+                    value=[],
+                    info="Select actions to process images, e.g., NoMonochromeAction" if self.config_manager.get_config("language") == "en" else "选择处理图片的动作，例如NoMonochromeAction"
+                )
+                action_param_components = gr.State(value={})
+                action_param_inputs = []
+                dataset_name_input = gr.Textbox(
+                    label=self.get_text("dataset_name"),
+                    value="waifuc_dataset",
+                    info="Name for the output dataset" if self.config_manager.get_config("language") == "en" else "输出数据集的名称"
+                )
+                exporter_dropdown = gr.Dropdown(
+                    choices=self.exporter_manager.exporter_names,
+                    label=self.get_text("select_exporter"),
+                    value=self.exporter_manager.exporter_names[0] if self.exporter_manager.exporter_names else None,
+                    info="Choose an exporter like TextualInversionExporter" if self.config_manager.get_config("language") == "en" else "选择导出器，如TextualInversionExporter"
+                )
+                start_btn = gr.Button(self.get_text("start_collection"))
+                status = gr.Textbox(label=self.get_text("status"), interactive=False)
+                output_file = gr.File(label=self.get_text("download_data"))
+            with gr.Tab(self.get_text("config")):
+                gr.Markdown("### " + self.get_text("config_management"))
+                pixiv_token_input = gr.Textbox(
+                    label=self.get_text("pixiv_token"),
+                    value=self.config_manager.get_config("pixiv_refresh_token"),
+                    type="password",
+                    info="Required for PixivSearchSource, stored in session only" if self.config_manager.get_config("language") == "en" else "PixivSearchSource所需，仅存储在会话中"
+                )
+                output_dir_input = gr.Textbox(
+                    label=self.get_text("output_dir"),
+                    value=self.config_manager.get_config("output_dir"),
+                    info="Directory for output data, session-specific" if self.config_manager.get_config("language") == "en" else "输出数据的目录，会话专用"
+                )
+                num_items_input = gr.Number(
+                    label=self.get_text("num_items"),
+                    value=self.config_manager.get_config("default_num_items"),
+                    info="Default number of images to collect" if self.config_manager.get_config("language") == "en" else "默认收集的图片数量"
+                )
+                resize_size_input = gr.Number(
+                    label=self.get_text("resize_size"),
+                    value=self.config_manager.get_config("default_resize_size"),
+                    info="Default size for ResizeAction" if self.config_manager.get_config("language") == "en" else "ResizeAction的默认尺寸"
+                )
+                config_export_btn = gr.Button(self.get_text("export_config"))
+                config_import_file = gr.File(label=self.get_text("import_config"))
+                save_config_btn = gr.Button(self.get_text("save_config"))
+                config_status = gr.Textbox(label=self.get_text("config_status"), interactive=False)
+            with gr.Tab(self.get_text("logs")):
+                gr.Markdown("### " + self.get_text("view_logs"))
+                log_output = gr.Textbox(label=self.get_text("view_logs"), interactive=False, lines=10)
+                log_download_btn = gr.Button(self.get_text("download_log"))
+                log_download_file = gr.File(label=self.get_text("download_log"))
+            def update_language(language):
+                self.config_manager.set_config("language", language)
+                return {
+                    "demo": demo.update(title=self.get_text("title")),
+                    "language_dropdown": language_dropdown.update(label=self.get_text("language")),
+                    "source_dropdown": source_dropdown.update(label=self.get_text("select_source")),
+                    "action_checkboxes": action_checkboxes.update(label=self.get_text("select_actions")),
+                    "dataset_name_input": dataset_name_input.update(label=self.get_text("dataset_name")),
+                    "exporter_dropdown": exporter_dropdown.update(label=self.get_text("select_exporter")),
+                    "start_btn": start_btn.update(value=self.get_text("start_collection")),
+                    "status": status.update(label=self.get_text("status")),
+                    "output_file": output_file.update(label=self.get_text("download_data")),
+                    "pixiv_token_input": pixiv_token_input.update(label=self.get_text("pixiv_token")),
+                    "output_dir_input": output_dir_input.update(label=self.get_text("output_dir")),
+                    "num_items_input": num_items_input.update(label=self.get_text("num_items")),
+                    "resize_size_input": resize_size_input.update(label=self.get_text("resize_size")),
+                    "config_export_btn": config_export_btn.update(value=self.get_text("export_config")),
+                    "config_import_file": config_import_file.update(label=self.get_text("import_config")),
+                    "save_config_btn": save_config_btn.update(value=self.get_text("save_config")),
+                    "config_status": config_status.update(label=self.get_text("config_status")),
+                    "log_output": log_output.update(label=self.get_text("view_logs")),
+                    "log_download_btn": log_download_btn.update(value=self.get_text("download_log")),
+                    "log_download_file": log_download_file.update(label=self.get_text("download_log"))
+                }
+            language_dropdown.change(
+                fn=update_language,
+                inputs=language_dropdown,
+                outputs=[
+                    demo, language_dropdown, source_dropdown, action_checkboxes,
+                    dataset_name_input, exporter_dropdown, start_btn, status,
+                    output_file, pixiv_token_input, output_dir_input, num_items_input,
+                    resize_size_input, config_export_btn, config_import_file,
+                    save_config_btn, config_status, log_output, log_download_btn,
+                    log_download_file
+                ]
+            )
+            source_dropdown.change(
+                fn=self.update_params_visibility,
+                inputs=source_dropdown,
+                outputs=param_components + [local_source_info]
+            )
+            action_checkboxes.change(
+                fn=self.update_action_params,
+                inputs=action_checkboxes,
+                outputs=[gr.State(), gr.State()]
+            )
+            def save_configs(pixiv_token, output_dir, num_items, resize_size):
+                self.config_manager.set_config("pixiv_refresh_token", pixiv_token)
+                self.config_manager.set_config("output_dir", output_dir)
+                self.config_manager.set_config("default_num_items", int(num_items))
+                self.config_manager.set_config("default_resize_size", int(resize_size))
+                return self.get_text("config_status") + ": " + ("Configuration saved" if self.config_manager.get_config("language") == "en" else "配置已保存")
+            save_config_btn.click(
+                fn=save_configs,
+                inputs=[pixiv_token_input, output_dir_input, num_items_input, resize_size_input],
+                outputs=config_status
+            )
+            config_export_btn.click(
+                fn=self.config_manager.export_config,
+                inputs=[],
+                outputs=gr.File(label=self.get_text("export_config"))
+            )
+            config_import_file.upload(
+                fn=self.config_manager.import_config,
+                inputs=config_import_file,
+                outputs=config_status
+            )
+            return {
+                "demo": demo,
+                "source_dropdown": source_dropdown,
+                "param_components": param_components,
+                "action_checkboxes": action_checkboxes,
+                "action_param_components": action_param_components,
+                "dataset_name_input": dataset_name_input,
+                "exporter_dropdown": exporter_dropdown,
+                "start_btn": start_btn,
+                "status": status,
+                "output_file": output_file,
+                "log_output": log_output,
+                "log_download_btn": log_download_btn,
+                "log_download_file": log_download_file
+            }