Upload 7 files

waifuc_gui/action_manager.py

@@ -0,0 +1,50 @@
+import inspect
+import waifuc.action
+import gradio as gr
+from .config_manager import ConfigManager
+
+class ActionManager:
+    def __init__(self, config_manager: ConfigManager):
+        self.config_manager = config_manager
+        self.action_classes = [
+            cls for name, cls in inspect.getmembers(waifuc.action, inspect.isclass)
+            if name.endswith("Action") and not inspect.isabstract(cls)
+        ]
+        self.action_names = [cls.__name__ for cls in self.action_classes]
+
+    def get_action_params(self, action_name):
+        action_cls = next(cls for cls in self.action_classes if cls.__name__ == action_name)
+        sig = inspect.signature(action_cls.__init__)
+        return [p for p in sig.parameters.values() if p.name != 'self']
+
+    def create_action_param_inputs(self, selected_actions):
+        params = {}
+        for action in selected_actions:
+            params[action] = {}
+            for param in self.get_action_params(action):
+                if param.name == "size" and action == "ResizeAction":
+                    params[action][param.name] = gr.Number(
+                        label="Resize Size" if self.config_manager.get_config("language") == "en" else "调整大小",
+                        value=self.config_manager.get_config("default_resize_size"),
+                        info="Image resize dimension" if self.config_manager.get_config("language") == "en" else "图片调整尺寸"
+                    )
+                elif param.annotation == int:
+                    params[action][param.name] = gr.Number(label=param.name)
+                else:
+                    params[action][param.name] = gr.Textbox(label=param.name)
+        return params
+
+    def instantiate_actions(self, selected_actions, action_params):
+        actions = []
+        for name in selected_actions:
+            action_cls = next(cls for cls in self.action_classes if cls.__name__ == name)
+            params = action_params.get(name, {})
+            if name == "FaceCountAction":
+                actions.append(action_cls(params.get("count", 1)))
+            elif name == "AlignMinSizeAction":
+                actions.append(action_cls(params.get("size", 800)))
+            elif name == "ResizeAction":
+                actions.append(action_cls(size=params.get("size", self.config_manager.get_config("default_resize_size"))))
+            else:
+                actions.append(action_cls(**params))
+        return actions

waifuc_gui/config_manager.py

@@ -0,0 +1,64 @@
+import os
+import json
+import uuid
+from typing import Dict, Any
+
+class ConfigManager:
+    def __init__(self, session_id: str):
+        self.session_id = session_id
+        self.config_dir = f"/tmp/user_{session_id}"
+        self.config_file = os.path.join(self.config_dir, "config.json")
+        self.config: Dict[str, Any] = {}
+        self.load_configs()
+
+    def load_configs(self):
+        self.config = {
+            "pixiv_refresh_token": "",
+            "output_dir": f"/tmp/user_{self.session_id}/output",
+            "default_num_items": 100,
+            "default_resize_size": 512,
+            "language": "zh"
+        }
+        if os.path.exists(self.config_file):
+            try:
+                with open(self.config_file, 'r') as f:
+                    self.config.update(json.load(f))
+            except Exception as e:
+                print(f"Failed to load session config: {e}")
+
+    def save_configs(self):
+        os.makedirs(self.config_dir, exist_ok=True)
+        try:
+            with open(self.config_file, 'w') as f:
+                json.dump(self.config, f, indent=2)
+        except Exception as e:
+            print(f"Failed to save session config: {e}")
+
+    def get_config(self, key: str, default: Any = None) -> Any:
+        return self.config.get(key, default)
+
+    def set_config(self, key: str, value: Any):
+        self.config[key] = value
+        self.save_configs()
+
+    def get_all_configs(self) -> Dict[str, Any]:
+        return self.config.copy()
+
+    def export_config(self) -> str:
+        export_path = os.path.join(self.config_dir, "config_export.json")
+        with open(export_path, 'w') as f:
+            json.dump(self.config, f, indent=2)
+        return export_path
+
+    def import_config(self, config_file):
+        with open(config_file, 'r') as f:
+            new_config = json.load(f)
+        self.config.update(new_config)
+        self.save_configs()
+        return "Configuration imported" if self.config.get("language") == "en" else "配置已导入"
+
+    def cleanup(self):
+        try:
+            shutil.rmtree(self.config_dir, ignore_errors=True)
+        except Exception as e:
+            print(f"Failed to cleanup session: {e}")

waifuc_gui/exporter_manager.py

@@ -0,0 +1,17 @@
+import inspect
+import waifuc.export
+from .config_manager import ConfigManager
+
+class ExporterManager:
+    def __init__(self, config_manager: ConfigManager):
+        self.config_manager = config_manager
+        self.exporter_classes = [
+            cls for name, cls in inspect.getmembers(waifuc.export, inspect.isclass)
+            if name.endswith("Exporter") and not inspect.isabstract(cls)
+        ]
+        self.exporter_names = [cls.__name__ for cls in self.exporter_classes]
+
+    def instantiate_exporter(self, selected_exporter, dataset_name):
+        exporter_cls = next(cls for cls in self.exporter_classes if cls.__name__ == selected_exporter)
+        output_dir = self.config_manager.get_config("output_dir")
+        return exporter_cls(dataset_name, path=output_dir)

waifuc_gui/file_handler.py

@@ -0,0 +1,29 @@
+import os
+import zipfile
+import shutil
+
+class FileHandler:
+    def extract_zip(self, uploaded_file):
+        extract_dir = "/tmp/input"
+        shutil.rmtree(extract_dir, ignore_errors=True)
+        os.makedirs(extract_dir)
+        with zipfile.ZipFile(uploaded_file, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+        return extract_dir
+
+    def create_zip(self, dataset_name):
+        output_dir = f"/tmp/{dataset_name}"
+        zip_path = f"/tmp/{dataset_name}.zip"
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, dirs, files in os.walk(dataset_name):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, '.')
+                    zipf.write(file_path, arcname)
+        return zip_path
+
+    def save_log(self, log_content):
+        log_path = "/tmp/waifuc_log.txt"
+        with open(log_path, 'w') as f:
+            f.write(log_content)
+        return log_path

waifuc_gui/interface.py

@@ -0,0 +1,244 @@
+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
+            }

waifuc_gui/source_manager.py

@@ -0,0 +1,86 @@
+import inspect
+import waifuc.source
+import gradio as gr
+from .config_manager import ConfigManager
+
+class SourceManager:
+    def __init__(self, config_manager: ConfigManager):
+        self.config_manager = config_manager
+        self.source_classes = [
+            cls for name, cls in inspect.getmembers(waifuc.source, inspect.isclass)
+            if name.endswith("Source") and not inspect.isabstract(cls)
+        ]
+        self.source_names = [cls.__name__ for cls in self.source_classes]
+
+    def get_source_params(self, selected_source):
+        source_cls = next(cls for cls in self.source_classes if cls.__name__ == selected_source)
+        sig = inspect.signature(source_cls.__init__)
+        return [p for p in sig.parameters.values() if p.name != 'self']
+
+    def create_param_inputs(self):
+        params = {}
+        for source in self.source_names:
+            if source == "LocalSource":
+                params[source] = {
+                    "file": gr.File(
+                        label="Upload ZIP file" if self.config_manager.get_config("language") == "en" else "上传zip文件",
+                        info="Upload a ZIP file containing images for LocalSource" if self.config_manager.get_config("language") == "en" else "上传包含图片的zip文件用于LocalSource",
+                        visible=False
+                    )
+                }
+            else:
+                source_cls = next(cls for cls in self.source_classes if cls.__name__ == source)
+                sig = inspect.signature(source_cls.__init__)
+                for param in sig.parameters.values():
+                    if param.name != 'self':
+                        if param.name == "refresh_token":
+                            params.setdefault(source, {})[param.name] = gr.Textbox(
+                                label="Pixiv Refresh Token" if self.config_manager.get_config("language") == "en" else "Pixiv刷新令牌",
+                                type="password",
+                                visible=False,
+                                value=self.config_manager.get_config("pixiv_refresh_token"),
+                                info="Set in Config tab or here" if self.config_manager.get_config("language") == "en" else "在配置选项卡或此处设置"
+                            )
+                        elif param.name == "num_items":
+                            params.setdefault(source, {})[param.name] = gr.Number(
+                                label="Number of Items" if self.config_manager.get_config("language") == "en" else "图片数量",
+                                visible=False,
+                                value=self.config_manager.get_config("default_num_items"),
+                                info="Number of images to collect" if self.config_manager.get_config("language") == "en" else "要收集的图片数量"
+                            )
+                        elif param.annotation == list:
+                            params.setdefault(source, {})[param.name] = gr.Textbox(
+                                label=param.name,
+                                placeholder="Comma-separated tags, e.g., tag1,tag2" if self.config_manager.get_config("language") == "en" else "逗号分隔的标签，例如：tag1,tag2",
+                                visible=False
+                            )
+                        elif param.annotation == int:
+                            params.setdefault(source, {})[param.name] = gr.Number(
+                                label=param.name,
+                                visible=False
+                            )
+                        else:
+                            params.setdefault(source, {})[param.name] = gr.Textbox(
+                                label=param.name,
+                                visible=False
+                            )
+        return params
+
+    def instantiate_source(self, selected_source, params, file_handler):
+        if selected_source == "LocalSource":
+            uploaded_file = params.get("file", None)
+            if not uploaded_file:
+                raise ValueError(
+                    "Please upload a ZIP file for LocalSource" if self.config_manager.get_config("language") == "en" else
+                    "请上传LocalSource的zip文件"
+                )
+            extract_dir = file_handler.extract_zip(uploaded_file)
+            return waifuc.source.LocalSource(extract_dir)
+        else:
+            source_cls = next(cls for cls in self.source_classes if cls.__name__ == selected_source)
+            source_params = {k: v for k, v in params.items() if k in [p.name for p in self.get_source_params(selected_source)]}
+            if "tags" in source_params:
+                source_params["tags"] = [tag.strip() for tag in source_params["tags"].split(',')]
+            if "refresh_token" in source_params and not source_params["refresh_token"]:
+                source_params["refresh_token"] = self.config_manager.get_config("pixiv_refresh_token")
+            return source_cls(**source_params)