{
"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.5 ~ | by ANXETY\n",
"\n",
"\n",
"\n",
"# =========================================================\n",
"import os\n",
"import sys\n",
"import time\n",
"import ipywidgets as widgets\n",
"from ipywidgets import Label, Button, VBox, HBox\n",
"from IPython.display import display, HTML\n",
"\n",
"\n",
"# ================= DETECT ENV =================\n",
"def detect_environment():\n",
" environments = {\n",
" 'COLAB_GPU': ('Google Colab', \"/content\"),\n",
" 'KAGGLE_URL_BASE': ('Kaggle', \"/kaggle/working/content\"),\n",
" 'SAGEMAKER_INTERNAL_IMAGE_URI': ('SageMaker Studio Lab', \"/home/studio-lab-user/content\")\n",
" }\n",
"\n",
" for env_var, (environment, path) in environments.items():\n",
" if env_var in os.environ:\n",
" return environment, path\n",
"\n",
" sys.exit(\"\\033[31mError: an unsupported runtime environment was detected.\\n\\033[34mSupported environments:\\033[0m Google Colab, Kaggle, Sagemaker Studio Lab\")\n",
"\n",
"env, root_path = detect_environment()\n",
"webui_path = f\"{root_path}/sdw\"\n",
"# ----------------------------------------------\n",
"\n",
"\n",
"directories = {\n",
" \"Images\": f\"{webui_path}/outputs\",\n",
" \"Models\": f\"{webui_path}/models/Stable-diffusion/\",\n",
" \"Vae\": f\"{webui_path}/models/VAE/\",\n",
" \"LoRa\": f\"{webui_path}/models/Lora/\",\n",
" \"ControlNet Models\": f\"{webui_path}/models/ControlNet/\"\n",
"}\n",
"\n",
"\n",
"# ==================== CSS ====================\n",
"CSS = \"\"\"\n",
"\n",
"\"\"\"\n",
"\n",
"display(HTML(CSS))\n",
"# ==================== CSS ====================\n",
"\n",
"\n",
"# ================ AutoCleaner function ================\n",
"def clean_directory(directory):\n",
" deleted_files = 0\n",
" for root, dirs, files in os.walk(directory):\n",
" for file in files:\n",
" if not file.endswith((\".txt\", \".yaml\")):\n",
" file_path = os.path.join(root, file)\n",
" os.remove(file_path)\n",
" deleted_files += 1\n",
"\n",
" return deleted_files\n",
"\n",
"\n",
"def get_word_variant(n, variants):\n",
" unit = abs(n) % 10\n",
" tens = abs(n) % 100\n",
"\n",
" if tens in range(11, 15):\n",
" return variants[2]\n",
" elif unit == 1:\n",
" return variants[0]\n",
" elif unit in range(2, 5):\n",
" return variants[1]\n",
" return variants[2]\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'{message}
')\n",
" display(message_widget)\n",
"\n",
"\n",
"def on_clear_button_press(button):\n",
" container.add_class(\"hide\")\n",
" time.sleep(0.3)\n",
" widgets.Widget.close_all()\n",
"\n",
"\n",
"def generate_messages(deleted_files_dict):\n",
" messages = []\n",
" word_variants = {\n",
" \"Images\": [\"Image\", \"Images\", \"Images\"],\n",
" \"Models\": [\"Model\", \"Models\", \"Models\"],\n",
" \"Vae\": [\"VAE\", \"VAE\", \"VAE\"],\n",
" \"LoRa\": [\"LoRa\", \"LoRa\", \"LoRa\"],\n",
" \"ControlNet Models\": [\"ControlNet Model\", \"ControlNet Models\", \"ControlNet Models\"]\n",
" }\n",
" deleted_word_variants = [\"Deleted\", \"Deleted\", \"Deleted\"]\n",
"\n",
" for key, value in deleted_files_dict.items():\n",
" word_variant = word_variants.get(key, [\"\", \"\", \"\"])\n",
" deleted_word = get_word_variant(value, deleted_word_variants)\n",
" object_word = get_word_variant(value, word_variant)\n",
"\n",
" messages.append(f\"{deleted_word} {value} {object_word}\")\n",
"\n",
" return messages\n",
"# ================ AutoCleaner function ================\n",
"\n",
"\n",
"# --- storage memory ---\n",
"import psutil\n",
"directory = os.getcwd()\n",
"disk_space = psutil.disk_usage(directory)\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",
"Use ctrl or shift for multiple selections.\n",
"''')\n",
"auto_cleaner_widget = widgets.SelectMultiple(options=AutoCleaner_options, layout=widgets.Layout(width='100%')).add_class(\"custom-select-multiple\")\n",
"output = widgets.Output().add_class(\"output\")\n",
"# ---\n",
"execute_button = Button(description='Execute Cleaning').add_class(\"button_execute\").add_class(\"button\")\n",
"execute_button.on_click(on_execute_button_press)\n",
"clear_button = Button(description='Hide Widget').add_class(\"button_clear\").add_class(\"button\")\n",
"clear_button.on_click(on_clear_button_press)\n",
"# ---\n",
"storage_info = widgets.HTML(f'''\n",
"Total storage: {total:.2f} GB | Used: {used:.2f} GB | Free: {free:.2f} GB
\n",
"''')\n",
"# ---\n",
"buttons = widgets.HBox([execute_button, clear_button])\n",
"lower_information_panel = widgets.HBox([buttons, storage_info]).add_class(\"lower_information_panel\")\n",
"\n",
"container = VBox([instruction_label, widgets.HTML('
'), auto_cleaner_widget, output, widgets.HTML('
'), lower_information_panel]).add_class(\"container\")\n",
"\n",
"display(container)"
],
"metadata": {
"id": "I22dFg7F2j3G"
},
"execution_count": null,
"outputs": []
}
]
}