{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "##~ AutoCleaner V3.7 CODE | BY: ANXETY   ~##\n",
        "\n",
        "from directory_setup import models_dir, vaes_dir, control_dir, loras_dir, output_dir\n",
        "\n",
        "import os\n",
        "import time\n",
        "import ipywidgets as widgets\n",
        "from ipywidgets import Label, Button, VBox, HBox\n",
        "from IPython.display import display, HTML, Javascript\n",
        "\n",
        "\n",
        "# Setup Env\n",
        "env = os.getenv('ENV_NAME')\n",
        "root_path = os.getenv('ROOT_PATH')\n",
        "webui_path = os.getenv('WEBUI_PATH')\n",
        "free_plan = os.getenv('FREE_PLAN')\n",
        "\n",
        "\n",
        "# ==================== CSS ====================\n",
        "# Main CSS\n",
        "css_file_path = f\"{root_path}/CSS/auto_cleaner.css\"\n",
        "with open(css_file_path , \"r\") as f:\n",
        "    CSS_AC = f.read()\n",
        "display(HTML(f\"<style>{CSS_AC}</style>\"))\n",
        "\n",
        "\n",
        "# ================ AutoCleaner function ================\n",
        "directories = {\n",
        "    \"Images\": output_dir,\n",
        "    \"Models\": models_dir,\n",
        "    \"Vae\": vaes_dir,\n",
        "    \"LoRa\": loras_dir,\n",
        "    \"ControlNet Models\": control_dir\n",
        "}\n",
        "\n",
        "\"\"\" functions \"\"\"\n",
        "def clean_directory(directory):\n",
        "    deleted_files = 0\n",
        "    image_dir = directories['Images']\n",
        "\n",
        "    for root, dirs, files in os.walk(directory):\n",
        "        for file in files:\n",
        "            file_path = os.path.join(root, file)\n",
        "\n",
        "            if file.endswith(\".txt\"):\n",
        "                continue\n",
        "            if file.endswith((\".safetensors\", \".pt\")) or root == image_dir: # fix for image counter\n",
        "                deleted_files += 1\n",
        "\n",
        "            os.remove(file_path)\n",
        "    return deleted_files\n",
        "\n",
        "def update_memory_info():\n",
        "    disk_space = psutil.disk_usage(os.getcwd())\n",
        "    total = disk_space.total / (1024 ** 3)\n",
        "    used = disk_space.used / (1024 ** 3)\n",
        "    free = disk_space.free / (1024 ** 3)\n",
        "\n",
        "    storage_info.value = f'''\n",
        "    <div class=\"storage_info_AC\">Total storage: {total:.2f} GB <span style=\"color: #555\">|</span> Used: {used:.2f} GB <span style=\"color: #555\">|</span> Free: {free:.2f} GB</div>\n",
        "    '''\n",
        "\n",
        "def on_execute_button_press(button):\n",
        "    selected_cleaners = auto_cleaner_widget.value\n",
        "    deleted_files_dict = {}\n",
        "\n",
        "    for option in selected_cleaners:\n",
        "        if option in directories:\n",
        "            deleted_files_dict[option] = clean_directory(directories[option])\n",
        "\n",
        "    output.clear_output()\n",
        "\n",
        "    with output:\n",
        "        for message in generate_messages(deleted_files_dict):\n",
        "            message_widget = HTML(f'<p class=\"output_message_AC\">{message}</p>')\n",
        "            display(message_widget)\n",
        "\n",
        "    update_memory_info()\n",
        "\n",
        "def on_clear_button_press(button):\n",
        "    container.add_class(\"hide\")\n",
        "    time.sleep(0.5)\n",
        "    widgets.Widget.close_all()\n",
        "\n",
        "def generate_messages(deleted_files_dict):\n",
        "    messages = []\n",
        "    word_variants = {\n",
        "        \"Images\": \"Images\",\n",
        "        \"Models\": \"Models\",\n",
        "        \"Vae\": \"Vae\",\n",
        "        \"LoRa\": \"LoRa\",\n",
        "        \"ControlNet Models\": \"ControlNet Models\"\n",
        "    }\n",
        "    for key, value in deleted_files_dict.items():\n",
        "        object_word = word_variants.get(key)\n",
        "        messages.append(f\"Deleted {value} {object_word}\")\n",
        "    return messages\n",
        "# ================ AutoCleaner function ================\n",
        "\n",
        "\n",
        "# --- storage memory ---\n",
        "import psutil\n",
        "disk_space = psutil.disk_usage(os.getcwd())\n",
        "total = disk_space.total / (1024 ** 3)\n",
        "used = disk_space.used / (1024 ** 3)\n",
        "free = disk_space.free / (1024 ** 3)\n",
        "\n",
        "\n",
        "# UI Code\n",
        "AutoCleaner_options = AutoCleaner_options = list(directories.keys())\n",
        "instruction_label = widgets.HTML('''\n",
        "<span class=\"instruction_AC\">Use <span style=\"color: #B2B2B2;\">ctrl</span> or <span style=\"color: #B2B2B2;\">shift</span> for multiple selections.</span>\n",
        "''')\n",
        "auto_cleaner_widget = widgets.SelectMultiple(options=AutoCleaner_options, layout=widgets.Layout(width='auto')).add_class(\"custom-select-multiple_AC\")\n",
        "output = widgets.Output().add_class(\"output_AC\")\n",
        "# ---\n",
        "execute_button = Button(description='Execute Cleaning').add_class(\"button_execute_AC\").add_class(\"button_AC\")\n",
        "execute_button.on_click(on_execute_button_press)\n",
        "clear_button = Button(description='Hide Widget').add_class(\"button_clear_AC\").add_class(\"button_AC\")\n",
        "clear_button.on_click(on_clear_button_press)\n",
        "# ---\n",
        "storage_info = widgets.HTML(f'''\n",
        "<div class=\"storage_info_AC\">Total storage: {total:.2f} GB <span style=\"color: #555\">|</span> Used: {used:.2f} GB <span style=\"color: #555\">|</span> Free: {free:.2f} GB</div>\n",
        "''')\n",
        "# ---\n",
        "buttons = widgets.HBox([execute_button, clear_button])\n",
        "lower_information_panel = widgets.HBox([buttons, storage_info]).add_class(\"lower_information_panel_AC\")\n",
        "\n",
        "container = VBox([instruction_label, widgets.HTML('<hr>'), auto_cleaner_widget, output, widgets.HTML('<hr>'), lower_information_panel]).add_class(\"container_AC\")\n",
        "\n",
        "display(container)"
      ],
      "metadata": {
        "id": "I22dFg7F2j3G"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}