Upload app.py

app.py

@@ -0,0 +1,107 @@
+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()