{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/panguin6010/kohya_ss_google_colab/blob/master/kohya_ss_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MvroZ9rJ1iqN"
      },
      "source": [
        "# Kohya SS WebUI Colab Setup\n",
        "\n",
        "This Colab workbook sets up a Kohya SS instance on Colab and provides a link to access the Kohya WebUI on Gradio Live. Kohya SS is a Python library that provides Stable Diffusion-based models for image, text, and audio generation tasks. This Colab workbook provides a convenient way for users to run Kohya SS without needing to install anything on their local machine.\n",
        "\n",
        "This workbook was inspired by the work of [Spaceginner](https://github.com/Spaceginner)'s original Colab workbook and the [Kohya SS project](https://github.com/bmaltais/kohya_ss) by [bmaltais](https://github.com/bmaltais). The Colab workbook was coded by [panguin6010](https://github.com/panguin6010) \n",
        "\n",
        "\n",
        "## Tutorials\n",
        "\n",
        "Before running this code, make sure you are familiar with using Colab workbooks and have a basic understanding of Kohya SS and its usage. You can find tutorials for these online. If you encounter any issues or have suggestions for improvement, feel free to contribute to the project.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DrAnm1um5vjh"
      },
      "source": [
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vmoRnFQEqOeY",
        "outputId": "09876c9a-d043-4881-d92f-6ed54313c390"
      },
      "outputs": [],
      "source": [
        "#@markdown #Step 1: Mounting Google Drive\n",
        "\n",
        "#@markdown The first step in setting up Kohya SS on Colab is to mount your Google Drive to the Colab notebook. This allows you to save and access files from your Google Drive in the Colab notebook.\n",
        "\n",
        "#@markdown To mount your Google Drive, run the following code block, which mounts your Google Drive to the /content/gdrive directory in the Colab notebook.\n",
        "\n",
        "\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mvQwnr4354BM"
      },
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "7ca7f6f727da46ac9a1149e69c16c81f",
            "77e5e07552b641cf9c368fb3939cb1d1",
            "235e01b92646444387ebd31ab945358e"
          ]
        },
        "id": "jnhm7ycMrLWb",
        "outputId": "63ba39ed-90c6-4b2d-f03e-61775587b083"
      },
      "outputs": [],
      "source": [
        "#@markdown #Kohya SS WebUI Installation\n",
        "\n",
        "#@markdown Now that your Google Drive is linked, we need to install the Kohya SS WebUI.\n",
        "\n",
        "#@markdown The code clones the [Kohya SS Google Colab](\"https://github.com/panguin6010/kohya_ss_google_colab\") repository and creates the necessary directories for Kohya SS to run. It then resets the git repository and pulls the latest changes. Finally, it displays a success message.\n",
        "\n",
        "#@markdown Note: If Google Drive is not connected, the code will use Colab storage instead.\n",
        "\n",
        "#@title\n",
        "# Import necessary libraries\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "from subprocess import getoutput\n",
        "import ipywidgets as widgets\n",
        "import sys\n",
        "import fileinput\n",
        "import os\n",
        "import time\n",
        "\n",
        "# WebUI Installation\n",
        "\n",
        "# Check if Google Drive is connected\n",
        "if not os.path.exists(\"/content/gdrive/MyDrive/\"):\n",
        "    print(\"Gdrive not connected, using colab storage ...\")\n",
        "    time.sleep(4)\n",
        "    !mkdir -p /content/gdrive/MyDrive/\n",
        "\n",
        "# Clone the repository and create necessary directories\n",
        "with capture.capture_output() as cap:\n",
        "    def inf(msg, style, wdth):\n",
        "        inf = widgets.Button(description=msg, disabled=True, button_style=style, layout=widgets.Layout(min_width=wdth))\n",
        "        display(inf)\n",
        "\n",
        "    %mkdir -p /content/gdrive/MyDrive/sd\n",
        "    %cd /content/gdrive/MyDrive/sd\n",
        "    !git clone https://github.com/panguin6010/kohya_ss_google_colab kohya_ss_colab\n",
        "    !mkdir -p /content/gdrive/MyDrive/sd/kohya_ss_colab/cache/huggingface\n",
        "    !ln -s /content/gdrive/MyDrive/sd/kohya_ss_colab/cache/huggingface /root/.cache/\n",
        "\n",
        "# Reset the git repository and pull the latest changes\n",
        "with capture.capture_output() as cap:\n",
        "    %cd /content/gdrive/MyDrive/sd/kohya_ss_colab/\n",
        "    !git reset --hard\n",
        "    time.sleep(1)\n",
        "\n",
        "print(\"Updating the repository...\")\n",
        "!git pull\n",
        "\n",
        "# Clear the output and display the success message\n",
        "clear_output()\n",
        "inf(\"✓ Done\", \"success\", \"50px\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8SrMhmFz7Lt4"
      },
      "source": [
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "54e929bcb37e4997a696d0becdecfd84",
            "43fbca3abb04401296967f819680f94f",
            "6d87b2c916394932b1a53382fe3cdb4e"
          ]
        },
        "id": "yjvkHRlDtDmV",
        "outputId": "06e1e873-b1ed-4403-c9a4-19ac1caa961b"
      },
      "outputs": [],
      "source": [
        "#@markdown #Requirements Installation\n",
        "\n",
        "#@markdown Now that we have downloaded the Kohya SS WebUI, we need to install the necessary requirements.\n",
        "\n",
        "# Print the status message\n",
        "print(\"Installing requirements...\")\n",
        "\n",
        "# Change the working directory to the project folder\n",
        "%cd /content/gdrive/MyDrive/sd/kohya_ss_colab/\n",
        "\n",
        "# Install the requirements\n",
        "with capture.capture_output() as cap:\n",
        "    # Uncomment the following line if you need to install specific versions of torch and torchvision\n",
        "    # !pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116\n",
        "    \n",
        "    # Install the requirements from the requirements.txt file\n",
        "    !pip install -r requirements.txt\n",
        "\n",
        "# Clear the output to keep the notebook clean\n",
        "clear_output()\n",
        "\n",
        "# Print the success message\n",
        "inf(\"✓ Done\", \"success\", \"50px\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FLDvlHm1tYud"
      },
      "source": [
        "\n",
        "---\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IzS3hvuTtTqW",
        "outputId": "9e629e1f-c8eb-43a2-9639-2583937ba81a"
      },
      "outputs": [],
      "source": [
        "#@markdown # Start Kohya ss WebUI\n",
        "\n",
        "User = \"\" #@param {type:\"string\"}\n",
        "Password = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown - Adding a username and password is not necessary but it will improve the security of your Kohya instance.\n",
        "#@markdown ______\n",
        "#@markdown # Please click the link that concludes with ```gradio.live``` to access your instance\n",
        "# Encourage users to contribute improvements\n",
        "print(\"Please feel free to make any changes or improvements you think would enhance this setup. Your input and contributions are greatly appreciated!\")\n",
        "# Check if the user has provided a username and password\n",
        "if User and Password:\n",
        "    # Run the Kohya GUI with the provided credentials\n",
        "    !python /content/gdrive/MyDrive/sd/kohya_ss_colab/kohya_gui.py --username $User --password $Password --share \n",
        "else:\n",
        "    # Run the Kohya GUI without credentials\n",
        "    !python /content/gdrive/MyDrive/sd/kohya_ss_colab/kohya_gui.py --share \n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "authorship_tag": "ABX9TyOZmOjfS55zOBmbTmRNOf3b",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "235e01b92646444387ebd31ab945358e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "43fbca3abb04401296967f819680f94f": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": "50px",
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "54e929bcb37e4997a696d0becdecfd84": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "success",
            "description": "✓ Done",
            "disabled": true,
            "icon": "",
            "layout": "IPY_MODEL_43fbca3abb04401296967f819680f94f",
            "style": "IPY_MODEL_6d87b2c916394932b1a53382fe3cdb4e",
            "tooltip": ""
          }
        },
        "6d87b2c916394932b1a53382fe3cdb4e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "77e5e07552b641cf9c368fb3939cb1d1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": "50px",
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7ca7f6f727da46ac9a1149e69c16c81f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "success",
            "description": "✓ Done",
            "disabled": true,
            "icon": "",
            "layout": "IPY_MODEL_77e5e07552b641cf9c368fb3939cb1d1",
            "style": "IPY_MODEL_235e01b92646444387ebd31ab945358e",
            "tooltip": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}