fast_dataset_maker.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n1tJPH73XWDk"
      },
      "source": [
        "[![Visitors](https://api.visitorbadge.io/api/combined?path=wsh.fast_dataset_maker&countColor=%232ccce4&style=flat&labelStyle=none)](https://visitorbadge.io/status?path=wsh.fast_dataset_maker)\n",
        "\n",
        "### SD-Lora-Training Colab\n",
        "| Notebook Name | Description | Link | Old-Version |\n",
        "| --- | --- | --- | --- |\n",
        "| [Colab_Lora_train](https://github.com/WSH032/lora-scripts/) | 基于[Akegarasu/lora-scripts](https://github.com/Akegarasu/lora-scripts)的定制化Colab notebook | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/WSH032/lora-scripts/blob/main/Colab_Lora_train.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/drive/1_f0qJdM43BSssNJWtgjIlk9DkIzLPadx) |\n",
        "| [kohya_train_webui](https://github.com/WSH032/kohya-config-webui) `NEW` | 基于[WSH032/kohya-config-webui](https://github.com/WSH032/kohya-config-webui)的WebUI版Colab notebook | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/WSH032/kohya-config-webui/blob/main/kohya_train_webui.ipynb) |\n",
        "\n",
        "---\n",
        "\n",
        "### 图片数据集 Colab\n",
        "| Notebook Name | Description | Link | Old-Version |\n",
        "| --- | --- | --- | --- |\n",
        "| [fast_dataset_maker](https://github.com/WSH032/sd-webui-fast-dataset-maker) `NEW` | 整合了 `图库`，`下载`，`去重`，`聚类` 的图片数据集WebUI | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/WSH032/sd-webui-fast-dataset-maker/blob/main/fast_dataset_maker.ipynb) |\n",
        "\n",
        "---\n",
        "\n",
        "### Development\n",
        "如果你觉得此项目有用，可以去 [![GitHub Repo stars](https://img.shields.io/github/stars/WSH032/sd-webui-fast-dataset-maker?style=social)](https://github.com/WSH032/sd-webui-fast-dataset-maker)</a>  点一颗小星星，非常感谢你⭐\n",
        "\n",
        "遇到问题可以在[Github上提issue ❓](https://github.com/WSH032/sd-webui-fast-dataset-maker/issues)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gvG7CCfHNnUu"
      },
      "source": [
        "# **A. Prepare env | 环境配置**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lJnNe99YlC5V"
      },
      "outputs": [],
      "source": [
        "####################\n",
        "\n",
        "import os\n",
        "import base64\n",
        "import sys\n",
        "\n",
        "from google.colab import output\n",
        "\n",
        "\n",
        "ROOT = \"/content\"\n",
        "rep_download_name = \"fast_dataset_maker\"\n",
        "encoded_data = b'aHR0cHM6Ly9naXRodWIuY29tL1dTSDAzMi9zZC13ZWJ1aS1mYXN0LWRhdGFzZXQtbWFrZXIuZ2l0'\n",
        "rep_url = base64.b64decode(encoded_data).decode('utf-8')\n",
        "rep_dir = os.path.join(ROOT, rep_download_name)\n",
        "extensions_dir_name = \"extensions\"\n",
        "extensions_dir = os.path.join(rep_dir, extensions_dir_name)\n",
        "requirments_versions_txt_path = os.path.join(rep_dir, \"requirements_versions.txt\")\n",
        "requirments_txt_path = os.path.join(rep_dir, \"requirements.txt\")\n",
        "deduplicate_cluster_dir = os.path.join(\n",
        "    extensions_dir,\n",
        "    \"image_deduplicate_cluster_webui\"\n",
        ")\n",
        "\n",
        "\n",
        "####################\n",
        "\n",
        "#@title ## 📦 1.1 Clone and Install | 克隆仓库和安装依赖\n",
        "\n",
        "#@markdown **此代码块运行完后请按键 `Ctrl + M + >` 重启内核 ！！！**\\\n",
        "#@markdown **此代码块运行完后请按键 `Ctrl + M + >` 重启内核 ！！！**\\\n",
        "#@markdown **此代码块运行完后请按键 `Ctrl + M + >` 重启内核 ！！！**\n",
        "\n",
        "#@markdown `代码执行程序` -> `重新启动代码执行程序`\n",
        "\n",
        "#@markdown **否则会出现 `AttributeError: module 'PIL.Image' has no attribute 'Resampling'`**\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown 强制指定依赖版本，勾选此项可能会修复依赖报错和潜在错误\n",
        "use_requirments_versions = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown 安静安装，会自动清除安装信息输出，如果需要debug请关闭\n",
        "quiet = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown 仓库分支\n",
        "branch = \"v1.1.0\" # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "def prepare_all_env(requirments_txt_path: str):\n",
        "    %pip install -r {requirments_txt_path}\n",
        "\n",
        "def clone():\n",
        "    print(f\"Clone from: {rep_url}\\n\")\n",
        "    !git clone {rep_url} --branch {branch} --recurse-submodules {rep_dir}\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "\n",
        "    clone()\n",
        "\n",
        "    if use_requirments_versions:\n",
        "        prepare_all_env(requirments_versions_txt_path)\n",
        "    else:\n",
        "        prepare_all_env(requirments_txt_path)\n",
        "\n",
        "    if quiet:\n",
        "        output.clear()\n",
        "\n",
        "    print(\"#\"*20)\n",
        "    print(\"\\n!!! Please Restart Kernel !!!\")\n",
        "    print(\"Ctrl + M + >\\n\")\n",
        "    print(\"#\"*20)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "AQXnh7yKsMi6"
      },
      "outputs": [],
      "source": [
        "####################\n",
        "\n",
        "import os\n",
        "import base64\n",
        "import sys\n",
        "\n",
        "from google.colab import output\n",
        "\n",
        "\n",
        "ROOT = \"/content\"\n",
        "rep_download_name = \"fast_dataset_maker\"\n",
        "encoded_data = b'aHR0cHM6Ly9naXRodWIuY29tL1dTSDAzMi9zZC13ZWJ1aS1mYXN0LWRhdGFzZXQtbWFrZXIuZ2l0'\n",
        "rep_url = base64.b64decode(encoded_data).decode('utf-8')\n",
        "rep_dir = os.path.join(ROOT, rep_download_name)\n",
        "extensions_dir_name = \"extensions\"\n",
        "extensions_dir = os.path.join(rep_dir, extensions_dir_name)\n",
        "requirments_versions_txt_path = os.path.join(rep_dir, \"requirements_versions.txt\")\n",
        "requirments_txt_path = os.path.join(rep_dir, \"requirements.txt\")\n",
        "deduplicate_cluster_dir = os.path.join(\n",
        "    extensions_dir,\n",
        "    \"image_deduplicate_cluster_webui\"\n",
        ")\n",
        "\n",
        "\n",
        "########## 上面的常量声明，请从第一个代码块复制过来，以确保重启后重新初始化 ##########\n",
        "\n",
        "#@title ## 🛠️ 1.2 Set config | 配置环境\n",
        "\n",
        "\n",
        "#@markdown **只要重启过内核就要再次运行此代码块**\\\n",
        "#@markdown **只要重启过内核就要再次运行此代码块**\\\n",
        "#@markdown **只要重启过内核就要再次运行此代码块**\n",
        "\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown 强制重写config\n",
        "force_rewrite = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown **1.2.1 infinite_image_browsing设置**\n",
        "\n",
        "#@markdown `infinite_image_browsing` 中 `文生图`,`图生图` 对应的路径\n",
        "outdir_txt2img_samples = \"/content/images\" #@param {type:\"string\"}\n",
        "outdir_img2img_samples = \"/content/images\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **1.2.2 dataset_tag_editor设置**\n",
        "\n",
        "#@markdown `dataset_tag_editor` 中显示的图片分辨，低分辨率可以加快显示，网络较差时请选择`128 / 64 `\n",
        "max_resolution = 192 #@param {type:\"number\"}\n",
        "\n",
        "\n",
        "####################\n",
        "import json\n",
        "import logging\n",
        "import textwrap\n",
        "from typing import List\n",
        "\n",
        "os.chdir(ROOT)\n",
        "\n",
        "\n",
        "# 自检\n",
        "extensions_name_list = [\n",
        "    \"dataset_tag_editor_standalone\",\n",
        "    \"Gelbooru_API_Downloader\",\n",
        "    \"image_deduplicate_cluster_webui\",\n",
        "    \"sd_webui_infinite_image_browsing\",\n",
        "]\n",
        "\n",
        "def check_extensions_dir(extensions_name_list: List[str], extensions_dir: str):\n",
        "    if not extensions_dir:\n",
        "        return\n",
        "\n",
        "    exist_extensions_folder_name_list = os.listdir(extensions_dir)\n",
        "    not_exist_extensions_name_list = [\n",
        "        extension\n",
        "        for extension in extensions_name_list\n",
        "        if extension not in exist_extensions_folder_name_list\n",
        "    ]\n",
        "    if not_exist_extensions_name_list:\n",
        "        logging.warning(\n",
        "            textwrap.dedent(f\"\"\"\\\n",
        "                以下扩展不存在于{extensions_dir}\\n{not_exist_extensions_name_list}\n",
        "                请检查 1.1 中是否正常Clone， 或者提issue: {rep_url}\n",
        "            \"\"\")\n",
        "        )\n",
        "check_extensions_dir(extensions_name_list, extensions_dir)\n",
        "\n",
        "\n",
        "# 扩展rep路径\n",
        "if rep_dir not in sys.path:\n",
        "    sys.path = [rep_dir] + sys.path\n",
        "\n",
        "\n",
        "# 设置 tag_editor 配置\n",
        "def set_tag_editor_config(force_rewrite: bool):\n",
        "    tag_editor_config_dict =dict(\n",
        "        use_temp_files = True,\n",
        "        temp_directory = \"\",\n",
        "        cleanup_tmpdir = False,\n",
        "        image_columns = 6,\n",
        "        max_resolution = max_resolution,\n",
        "        filename_word_regex = \"\",\n",
        "        tagger_use_spaces = True,\n",
        "        interrogator_use_cpu = False,\n",
        "        interrogator_keep_in_memory = False,\n",
        "        interrogator_max_length = 30,\n",
        "        interrogator_model_dir = \"\",\n",
        "    )\n",
        "    tag_editor_config_path = os.path.join(\n",
        "        extensions_dir,\n",
        "        \"dataset_tag_editor_standalone\",\n",
        "        \"settings.json\",\n",
        "\n",
        "    )\n",
        "    # 用户可能会在WebUI里自行设置，如果没开强制重写，这里不要重置了\n",
        "    if not os.path.exists(tag_editor_config_path) or force_rewrite:\n",
        "        with open(tag_editor_config_path, \"w\") as f:\n",
        "            json.dump(tag_editor_config_dict, f)\n",
        "set_tag_editor_config(force_rewrite)\n",
        "\n",
        "\n",
        "# 设置 image_browsing_config\n",
        "sd_webui_config_path = os.path.join(\n",
        "    extensions_dir,\n",
        "    \"sd_webui_infinite_image_browsing\",\n",
        "    \"config.json\",\n",
        ")\n",
        "def set_image_browsing_config():\n",
        "    sd_webui_config_dict = dict(\n",
        "        outdir_txt2img_samples = outdir_txt2img_samples,\n",
        "        outdir_img2img_samples = outdir_img2img_samples,\n",
        "        # outdir_extras_samples = \"\",\n",
        "    )\n",
        "    # 用户可能会在WebUI里自行设置，如果没开强制重写，这里不要重置了\n",
        "    if not os.path.exists(sd_webui_config_path) or force_rewrite:\n",
        "        with open(sd_webui_config_path, \"w\") as f:\n",
        "            json.dump(sd_webui_config_dict, f)\n",
        "set_image_browsing_config()\n",
        "\n",
        "\n",
        "# 导入 downloader\n",
        "downloader_dir = os.path.join(\n",
        "    extensions_dir,\n",
        "    \"Gelbooru_API_Downloader\",\n",
        ")\n",
        "sys.path = [downloader_dir] + sys.path\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tG3VwOskOH4c"
      },
      "source": [
        "# **B. Fast-Dataset-Maker | 功能模块**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MBr3XE63D1Y-"
      },
      "outputs": [],
      "source": [
        "#@title ### ⏬ Download images | 下载图片\n",
        "\n",
        "#@markdown 你可以在这找到tags规则 | some useful info for tags\n",
        "\n",
        "#@markdown - [tags](https://gelbooru.com/index.php?page=wiki&s=&s=view&id=25921)\n",
        "\n",
        "#@markdown - [cheatsheet](https://gelbooru.com/index.php?page=wiki&s=&s=view&id=26263)\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown 符合gelbooru搜索规则的tags\n",
        "tags = \"1girl\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown 需要下载的 `图片数`, `下载路径`, `下载单位`\n",
        "max_images_number = 100 #@param {type:\"integer\"}\n",
        "download_dir = \"/content/images\" #@param {type:\"string\"}\n",
        "unit = 100 #@param {type:\"slider\", min:1, max:100, step:1}\n",
        "\n",
        "#@markdown 与图片一起被下载的txt中tag规则：`逗号分割`，`移除下划线`，`对括号进行转义`\n",
        "add_comma = True #@param {type:\"boolean\"}\n",
        "remove_underscore = True #@param {type:\"boolean\"}\n",
        "use_escape = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown 对下载的图片进行完整性检查，以避免出现截断图片\n",
        "\n",
        "#@markdown | -1 | 0 | 1 | 2 |\n",
        "#@markdown |------|------|------|------|\n",
        "#@markdown | 不检查 | 检查但只输出信息 | 检查并修复 | 检查并删除 |\n",
        "\n",
        "\n",
        "check_images_mode = -1 #@param {type:\"slider\", min:-1, max:2, step:1}\n",
        "\n",
        "\n",
        "from download_images_coroutine import Scrape_images\n",
        "\n",
        "check_images_mode = None if check_images_mode == -1 else 1\n",
        "\n",
        "Scrape_images_kwargs = dict(\n",
        "    tags = tags,\n",
        "    max_images_number = max_images_number,\n",
        "    max_workers = 10,\n",
        "    download_dir = download_dir,\n",
        "    unit = unit,\n",
        "    add_comma = add_comma,\n",
        "    remove_underscore = remove_underscore,\n",
        "    use_escape = use_escape,\n",
        "    check_images_mode = check_images_mode,\n",
        "    )\n",
        "await Scrape_images(**Scrape_images_kwargs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "m_nVPTpQvvF5"
      },
      "outputs": [],
      "source": [
        "#@title ### 🚂 Launch WebUI | 启动WebUI\n",
        "\n",
        "#@markdown **在使用查重和聚类的时候，开启`缓存功能`，可加快图片显示**\\\n",
        "#@markdown **在使用查重和聚类的时候，开启`缓存功能`，可加快图片显示**\\\n",
        "#@markdown **在使用查重和聚类的时候，开启`缓存功能`，可加快图片显示**\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown `infinite_image_browsing`额外路径，`\",\"`分割\n",
        "extra_paths_str = \"/content, /content/images\" # @param {type:\"string\"}\n",
        "\n",
        "#@markdown `debug模式`：允许在colab中显示输出，但是会阻塞程序\n",
        "debug = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown `queue队列模式`：\n",
        "#@markdown - 如果无法正常工作，请尝试开启队列\n",
        "#@markdown - 开启后可能速度会变慢\n",
        "queue = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown 在Colab中以`iframe`显示\n",
        "inline = False #@param {type:\"boolean\"}\n",
        "\n",
        "import portpicker\n",
        "from google.colab.output import serve_kernel_port_as_iframe, serve_kernel_port_as_window\n",
        "\n",
        "from webui import WebuiUtils\n",
        "\n",
        "os.chdir(ROOT)\n",
        "\n",
        "# 关闭上一个demo\n",
        "def clear_last_demo():\n",
        "    last_webui_utils = globals().get(\"webui_utils\", None)\n",
        "    if last_webui_utils is not None:\n",
        "        last_webui_utils.demo.close()\n",
        "clear_last_demo()\n",
        "\n",
        "# 找一个空闲端口\n",
        "server_port = portpicker.PickUnusedPort()\n",
        "\n",
        "extra_paths = list(map(str.strip, extra_paths_str.split(\",\")))\n",
        "# WebUI的启动参数\n",
        "WebuiUtils_kwargs = dict(\n",
        "    sd_webui_config = sd_webui_config_path,\n",
        "    extra_paths = extra_paths,\n",
        ")\n",
        "webui_utils = WebuiUtils(**WebuiUtils_kwargs)  # 实例化\n",
        "\n",
        "if queue:\n",
        "    webui_utils.queue(concurrency_count=10)\n",
        "webui_utils.launch(server_port=server_port, inline=inline, debug=debug, share=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcJKURV-jC-S"
      },
      "source": [
        "# **C. Utils | 实用工具**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gFoW3f5Go4B2"
      },
      "outputs": [],
      "source": [
        "#@title ### 🔗 Mount Google Drive | 挂载谷歌硬盘\n",
        "\n",
        "import logging\n",
        "\n",
        "def mount_google_drive():\n",
        "    from google.colab import drive\n",
        "    print(\"Connecting to Google Drive...\")\n",
        "    drive.mount('/content/drive')\n",
        "\n",
        "try:\n",
        "    mount_google_drive()\n",
        "    print(\"✅ Done\")\n",
        "except Exception as e:\n",
        "    logging.error(f\"Error when mounting google drive\\nerro: {e}\")\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SzxZmHL_jHxv"
      },
      "outputs": [],
      "source": [
        "#@title ### 💾 ZIP | 压缩\n",
        "\n",
        "need_zip_path = \"/content/images\" #@param {type:\"string\"}\n",
        "output_zip_path = \"/content/drive/MyDrive/fast_dataset_maker/my_zip.zip\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "from modules.utils import zip_utils\n",
        "\n",
        "zip_utils.zip(need_zip_path, output_zip_path)\n",
        "print(\"✅ Done\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MG1UJYVykpBv"
      },
      "outputs": [],
      "source": [
        "#@title ### 📂 UNZIP | 解压\n",
        "\n",
        "\n",
        "need_extract_path = \"/content/drive/MyDrive/fast_dataset_maker/my_zip.zip\" #@param {type:\"string\"}\n",
        "output_extract_path = \"/content/images\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "from modules.utils import zip_utils\n",
        "\n",
        "zip_utils.unzip(need_extract_path, output_extract_path)\n",
        "print(\"✅ Done\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AhYYpSG9OhiY"
      },
      "source": [
        "# **D. Help | 帮助**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lq1loFEeCaB9"
      },
      "outputs": [],
      "source": [
        "#@title ### 😇 `WebuiUtils` Help | 查看`WebuiUtils`允许参数\n",
        "WebuiUtils(help=True)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}