Update app.py

app.py

@@ -1,107 +1,103 @@
-import os
-import logging
-import uuid
-import shutil
-from waifuc_gui.source_manager import SourceManager
-from waifuc_gui.action_manager import ActionManager
-from waifuc_gui.exporter_manager import ExporterManager
-from waifuc_gui.file_handler import FileHandler
-from waifuc_gui.interface import Interface
-from waifuc_gui.config_manager import ConfigManager
-import gradio as gr
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("waifuc_gui")
-log_stream = []
-
-class LogHandler(logging.Handler):
-    def emit(self, record):
-        log_stream.append(self.format(record))
-
-log_handler = LogHandler()
-log_handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
-logger.addHandler(log_handler)
-
-def start_collection(
-    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:
-        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(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(log_stream)
-
-def download_log(file_handler):
-    log_content = "\n".join(log_stream)
-    return file_handler.save_log(log_content)
-
-def cleanup_session(config_manager):
-    config_manager.cleanup()
-    return "Session cleaned" if config_manager.get_config("language") == "en" else "会话已清理"
-
-def main():
-    session_id = str(uuid.uuid4())
-    config_manager = ConfigManager(session_id)
-    source_manager = SourceManager(config_manager)
-    action_manager = ActionManager(config_manager)
-    exporter_manager = ExporterManager(config_manager)
-    file_handler = FileHandler()
-    interface = Interface(source_manager, action_manager, exporter_manager, config_manager)
-
-    ui = interface.build()
-    ui["start_btn"].click(
-        fn=start_collection,
-        inputs=[
-            ui["source_dropdown"],
-            gr.State(lambda: interface.collect_params(
-                ui["source_dropdown"].value,
-                *[p.value for p in ui["param_components"]]
-            )),
-            ui["action_checkboxes"],
-            ui["action_param_components"],
-            ui["dataset_name_input"],
-            ui["exporter_dropdown"]
-        ],
-        outputs=[ui["status"], ui["output_file"], ui["log_output"]],
-        _js="() => {return []}",
-        api_name="start_collection"
-    )
-
-    ui["log_download_btn"].click(
-        fn=download_log,
-        inputs=[gr.State(value=file_handler)],
-        outputs=ui["log_download_file"]
-    )
-
-    cleanup_btn = gr.Button("Cleanup Session" if config_manager.get_config("language") == "en" else "清理会话")
-    cleanup_status = gr.Textbox(label="Cleanup Status" if config_manager.get_config("language") == "en" else "清理状态")
-    cleanup_btn.click(
-        fn=cleanup_session,
-        inputs=[gr.State(value=config_manager)],
-        outputs=cleanup_status
-    )
-
-    ui["demo"].launch()
-
-if __name__ == "__main__":
-    main()
+import os
+import logging
+import uuid
+import shutil
+from waifuc_gui.source_manager import SourceManager
+from waifuc_gui.action_manager import ActionManager
+from waifuc_gui.exporter_manager import ExporterManager
+from waifuc_gui.file_handler import FileHandler
+from waifuc_gui.interface import Interface
+from waifuc_gui.config_manager import ConfigManager
+import gradio as gr
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("waifuc_gui")
+log_stream = []
+
+class LogHandler(logging.Handler):
+    def emit(self, record):
+        log_stream.append(self.format(record))
+
+log_handler = LogHandler()
+log_handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
+logger.addHandler(log_handler)
+
+def start_collection(
+    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:
+        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(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(log_stream)
+
+def download_log(file_handler):
+    log_content = "\n".join(log_stream)
+    return file_handler.save_log(log_content)
+
+def cleanup_session(config_manager):
+    config_manager.cleanup()
+    return "Session cleaned" if config_manager.get_config("language") == "en" else "会话已清理"
+
+def main():
+    session_id = str(uuid.uuid4())
+    config_manager = ConfigManager(session_id)
+    source_manager = SourceManager(config_manager)
+    action_manager = ActionManager(config_manager)
+    exporter_manager = ExporterManager(config_manager)
+    file_handler = FileHandler()
+    interface = Interface(source_manager, action_manager, exporter_manager, config_manager)
+
+    ui = interface.build()
+    ui["start_btn"].click(
+        fn=start_collection,
+        inputs=[
+            ui["source_dropdown"],
+            gr.State(lambda: interface.collect_params(
+                ui["source_dropdown"].value,
+                *[p.value for p in ui["param_components"]]
+            )),
+            ui["action_checkboxes"],
+            ui["action_param_components"],
+            ui["dataset_name_input"],
+            ui["exporter_dropdown"]
+        ],
+        outputs=[ui["status"], ui["output_file"], ui["log_output"]],
+        api_name="start_collection"
+    )
+
+    ui["log_download_btn"].click(
+        fn=download_log,
+        inputs=[gr.State(value=file_handler)],
+        outputs=ui["log_download_file"]
+    )
+
+    cleanup_btn = gr.Button("Cleanup Session" if config_manager.get_config("language") == "en" else "清理会话")
+    cleanup_status = gr.Textbox(label="Cleanup Status" if config_manager.get_config("language") == "en" else "清理状态")
+    cleanup_btn.click(
+        fn=cleanup_session,
+        inputs=[gr.State(value=config_manager)],
+        outputs=cleanup_status
+    )
+
+    ui["demo"].launch()