Update waifuc_gui/interface.py

waifuc_gui/interface.py

@@ -93,6 +93,46 @@ class Interface:
                 components.append(param_input)
         return components, action_params
 
+    def start_collection(
+        self, selected_source, params, selected_actions, action_params, dataset_name, selected_exporter,
+        source_manager, action_manager, exporter_manager, file_handler
+    ):
+        output_dir = f"/tmp/user_{source_manager.config_manager.session_id}/{dataset_name}"
+        try:
+            output_dir = f"/tmp/user_{source_manager.config_manager.session_id}/{dataset_name}"
+            logger = logging.getLogger("waifuc_gui")
+            logger.info(f"Starting collection for {selected_source} with dataset {dataset_name}")
+            os.makedirs(output_dir, exist_ok=True)
+            source = source_manager.instantiate_source(selected_source, params, file_handler)
+            actions = action_manager.instantiate_actions(selected_actions, action_params)
+            exporter = exporter_manager.instantiate_exporter(selected_exporter, dataset_name)
+
+            logger.info(f"Attaching actions: {selected_actions}")
+            os.chdir(f"/tmp/user_{source_manager.config_manager.session_id}")
+            source.attach(*actions).export(exporter)
+            logger.info(f"Export completed, creating ZIP")
+            zip_path = file_handler.create_zip(dataset_name)
+            logger.info(f"Collection completed: {zip_path}")
+            return (
+                f"Data collection completed, output file: {dataset_name}.zip" if source_manager.config_manager.get_config("language") == "en" else 
+                f"数据收集完成，输出文件：{dataset_name}.zip"
+            ), zip_path, "\n".join(logger.handlers[0].log_stream)
+        except Exception as e:
+            logger.error(f"Collection failed: {str(e)}")
+            return (
+                f"Data collection failed: {str(e)}" if source_manager.config_manager.get_config("language") == "en" else 
+                f"数据收集失败：{str(e)}"
+            ), None, "\n".join(logger.handlers[0].log_stream)
+
+    def download_log(self, file_handler):
+        logger = logging.getLogger("waifuc_gui")
+        log_content = "\n".join(logger.handlers[0].log_stream)
+        return file_handler.save_log(log_content)
+
+    def cleanup_session(self, config_manager):
+        config_manager.cleanup()
+        return "Session cleaned" if config_manager.get_config("language") == "en" else "会话已清理"
+
     def build(self):
         with gr.Blocks(title=self.get_text("title")) as demo:
             language_dropdown = gr.Dropdown(
@@ -167,6 +207,8 @@ class Interface:
                 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 {
@@ -234,18 +276,53 @@ class Interface:
                 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
-            }
+            start_btn.click(
+                fn=self.start_collection,
+                inputs=[
+                    source_dropdown,
+                    gr.State(lambda: self.collect_params(
+                        source_dropdown.value,
+                        *[p.value for p in param_components]
+                    )),
+                    action_checkboxes,
+                    action_param_components,
+                    dataset_name_input,
+                    exporter_dropdown,
+                    gr.State(value=self.source_manager),
+                    gr.State(value=self.action_manager),
+                    gr.State(value=self.exporter_manager),
+                    gr.State(value=FileHandler())
+                ],
+                outputs=[status, output_file, log_output],
+                api_name="start_collection"
+            )
+            log_download_btn.click(
+                fn=self.download_log,
+                inputs=[gr.State(value=FileHandler())],
+                outputs=log_download_file
+            )
+            cleanup_btn = gr.Button("Cleanup Session" if self.config_manager.get_config("language") == "en" else "清理会话")
+            cleanup_status = gr.Textbox(label="Cleanup Status" if self.config_manager.get_config("language") == "en" else "清理状态")
+            cleanup_btn.click(
+                fn=self.cleanup_session,
+                inputs=[gr.State(value=self.config_manager)],
+                outputs=cleanup_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,
+            "cleanup_btn": cleanup_btn,
+            "cleanup_status": cleanup_status
+        }